feat: Add context managers

ecederstrand · Oct 31, 2023 · 37d6ba2 · 37d6ba2
1 parent 04d7807
commit 37d6ba2
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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