From 37d6ba2934eafe3b9a1c73c91558faec88d81d90 Mon Sep 17 00:00:00 2001
From: ecederstrand <erik@cederstrand.dk>
Date: Tue, 31 Oct 2023 18:36:10 +0100
Subject: [PATCH] feat: Add context managers

---
 exchangelib/account.py             | 10 ++++++++++
 exchangelib/folders/base.py        |  6 +++---
 exchangelib/folders/collections.py | 18 +++++++++---------
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/exchangelib/account.py b/exchangelib/account.py
index 61c5c54f..b26feb35 100644
--- a/exchangelib/account.py
+++ b/exchangelib/account.py
@@ -54,6 +54,7 @@
     ToDoSearch,
     VoiceMail,
 )
+from .folders.collections import PullSubscription, PushSubscription, StreamingSubscription
 from .items import ALL_OCCURRENCES, AUTO_RESOLVE, HARD_DELETE, ID_ONLY, SAVE_ONLY, SEND_TO_NONE
 from .properties import EWSElement, Mailbox, SendingAs
 from .protocol import Protocol
@@ -793,6 +794,15 @@ def subscribe_to_streaming(self, event_types=None):
             event_types = SubscribeToStreaming.EVENT_TYPES
         return SubscribeToStreaming(account=self).get(folders=None, event_types=event_types)
 
+    def pull_subscription(self, **kwargs):
+        return PullSubscription(target=self, **kwargs)
+
+    def push_subscription(self, **kwargs):
+        return PushSubscription(target=self, **kwargs)
+
+    def streaming_subscription(self, **kwargs):
+        return StreamingSubscription(target=self, **kwargs)
+
     def unsubscribe(self, subscription_id):
         """Unsubscribe. Only applies to pull and streaming notifications.
 
diff --git a/exchangelib/folders/base.py b/exchangelib/folders/base.py
index 5c2f3af7..5bccd33d 100644
--- a/exchangelib/folders/base.py
+++ b/exchangelib/folders/base.py
@@ -631,15 +631,15 @@ def subscribe_to_streaming(self, event_types=None):
 
     @require_id
     def pull_subscription(self, **kwargs):
-        return PullSubscription(folder=self, **kwargs)
+        return PullSubscription(target=self, **kwargs)
 
     @require_id
     def push_subscription(self, **kwargs):
-        return PushSubscription(folder=self, **kwargs)
+        return PushSubscription(target=self, **kwargs)
 
     @require_id
     def streaming_subscription(self, **kwargs):
-        return StreamingSubscription(folder=self, **kwargs)
+        return StreamingSubscription(target=self, **kwargs)
 
     def unsubscribe(self, subscription_id):
         """Unsubscribe. Only applies to pull and streaming notifications.
diff --git a/exchangelib/folders/collections.py b/exchangelib/folders/collections.py
index 81595a43..c91efba9 100644
--- a/exchangelib/folders/collections.py
+++ b/exchangelib/folders/collections.py
@@ -448,13 +448,13 @@ def subscribe_to_streaming(self, event_types=None):
         return SubscribeToStreaming(account=self.account).get(folders=self.folders, event_types=event_types)
 
     def pull_subscription(self, **kwargs):
-        return PullSubscription(folder=self, **kwargs)
+        return PullSubscription(target=self, **kwargs)
 
     def push_subscription(self, **kwargs):
-        return PushSubscription(folder=self, **kwargs)
+        return PushSubscription(target=self, **kwargs)
 
     def streaming_subscription(self, **kwargs):
-        return StreamingSubscription(folder=self, **kwargs)
+        return StreamingSubscription(target=self, **kwargs)
 
     def unsubscribe(self, subscription_id):
         """Unsubscribe. Only applies to pull and streaming notifications.
@@ -540,8 +540,8 @@ def sync_hierarchy(self, sync_state=None, only_fields=None):
 
 
 class BaseSubscription(metaclass=abc.ABCMeta):
-    def __init__(self, folder, **subscription_kwargs):
-        self.folder = folder
+    def __init__(self, target, **subscription_kwargs):
+        self.target = target
         self.subscription_kwargs = subscription_kwargs
         self.subscription_id = None
 
@@ -550,19 +550,19 @@ def __enter__(self):
         """Create the subscription"""
 
     def __exit__(self, *args, **kwargs):
-        self.folder.unsubscribe(subscription_id=self.subscription_id)
+        self.target.unsubscribe(subscription_id=self.subscription_id)
         self.subscription_id = None
 
 
 class PullSubscription(BaseSubscription):
     def __enter__(self):
-        self.subscription_id, watermark = self.folder.subscribe_to_pull(**self.subscription_kwargs)
+        self.subscription_id, watermark = self.target.subscribe_to_pull(**self.subscription_kwargs)
         return self.subscription_id, watermark
 
 
 class PushSubscription(BaseSubscription):
     def __enter__(self):
-        self.subscription_id, watermark = self.folder.subscribe_to_push(**self.subscription_kwargs)
+        self.subscription_id, watermark = self.target.subscribe_to_push(**self.subscription_kwargs)
         return self.subscription_id, watermark
 
     def __exit__(self, *args, **kwargs):
@@ -572,5 +572,5 @@ def __exit__(self, *args, **kwargs):
 
 class StreamingSubscription(BaseSubscription):
     def __enter__(self):
-        self.subscription_id = self.folder.subscribe_to_streaming(**self.subscription_kwargs)
+        self.subscription_id = self.target.subscribe_to_streaming(**self.subscription_kwargs)
         return self.subscription_id