Support for context managers in ZeroConf

shenek · shenek · commit 13b627c0fc34 · 2020-07-08T16:09:56.000+02:00
diff --git a/zeroconf/__init__.py b/zeroconf/__init__.py
@@ -35,7 +35,8 @@
 import time
 import warnings
 from collections import OrderedDict
-from typing import Dict, List, Optional, Sequence, Union, cast
+from types import TracebackType  # noqa # used in type hints
+from typing import Dict, List, Optional, Sequence, Type, Union, cast
 from typing import Any, Callable, Set, Tuple  # noqa # used in type hints
 
 import ifaddr
@@ -2861,3 +2862,15 @@ def close(self) -> None:
             self.reaper.join()
             for s in self._respond_sockets:
                 s.close()
+
+    def __enter__(self) -> 'Zeroconf':
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> 'Literal[False]':  # type: ignore # noqa
+        self.close()
+        return False
diff --git a/zeroconf/test.py b/zeroconf/test.py
@@ -515,6 +515,15 @@ def test_launch_and_close(self):
         rv = r.Zeroconf(interfaces=r.InterfaceChoice.Default)
         rv.close()
 
+    def test_launch_and_close_context_manager(self):
+        with r.Zeroconf(interfaces=r.InterfaceChoice.All) as rv:
+            assert rv.done is False
+        assert rv.done is True
+
+        with r.Zeroconf(interfaces=r.InterfaceChoice.Default) as rv:
+            assert rv.done is False
+        assert rv.done is True
+
     @unittest.skipIf(not socket.has_ipv6, 'Requires IPv6')
     @unittest.skipIf(os.environ.get('SKIP_IPV6'), 'IPv6 tests disabled')
     def test_launch_and_close_v4_v6(self):
