python-zeroconf · bdraco · Jun 12, 2021 · Jun 12, 2021
diff --git a/Makefile b/Makefile
@@ -29,14 +29,16 @@ flake8:
 	flake8 --max-line-length=$(MAX_LINE_LENGTH) setup.py examples zeroconf
 
 pylint:
-	pylint zeroconf/__init__.py zeroconf/aio.py zeroconf/asyncio.py
+	pylint zeroconf
 
 .PHONY: black_check
 black_check:
 	black --check setup.py examples zeroconf
 
 mypy:
-	mypy examples/*.py zeroconf/*.py
+# --no-warn-redundant-casts --no-warn-unused-ignores is needed since we support multiple python versions
+# We should be able to drop this once python 3.6 goes away
+	mypy --no-warn-redundant-casts --no-warn-unused-ignores examples/*.py zeroconf
 
 test:
 	pytest -v tests

diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
@@ -0,0 +1,21 @@
+""" Multicast DNS Service Discovery for Python, v0.14-wmcbrine
+    Copyright 2003 Paul Scott-Murphy, 2014 William McBrine
+
+    This module provides a framework for the use of DNS Service Discovery
+    using IP multicast.
+
+    This library is free software; you can redistribute it and/or
+    modify it under the terms of the GNU Lesser General Public
+    License as published by the Free Software Foundation; either
+    version 2.1 of the License, or (at your option) any later version.
+
+    This library is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+    Lesser General Public License for more details.
+
+    You should have received a copy of the GNU Lesser General Public
+    License along with this library; if not, write to the Free Software
+    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301
+    USA
+"""
diff --git a/tests/utils/test_aio.py b/tests/utils/test_aio.py
@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+
+"""Unit tests for zeroconf.utils.aio."""
+
+import asyncio
+
+import pytest
+
+from zeroconf.utils import aio as aioutils
+
+
+@pytest.mark.asyncio
+async def test_get_running_loop_from_async() -> None:
+    """Test we can get the event loop."""
+    assert isinstance(aioutils.get_running_loop(), asyncio.AbstractEventLoop)
+
+
+def test_get_running_loop_no_loop() -> None:
+    """Test we get None when there is no loop running."""
+    assert aioutils.get_running_loop() is None
diff --git a/zeroconf/aio.py b/zeroconf/aio.py
@@ -46,6 +46,7 @@
     instance_name_from_service_info,
     millis_to_seconds,
 )
+from .utils.aio import wait_condition_or_timeout
 
 
 def _get_best_available_queue() -> queue.Queue:
@@ -55,35 +56,6 @@ def _get_best_available_queue() -> queue.Queue:
     return queue.Queue()
 
 
-# Switch to asyncio.wait_for once https://bugs.python.org/issue39032 is fixed
-async def wait_condition_or_timeout(condition: asyncio.Condition, timeout: float) -> None:
-    """Wait for a condition or timeout."""
-    loop = asyncio.get_event_loop()
-    future = loop.create_future()
-
-    def _handle_timeout() -> None:
-        if not future.done():
-            future.set_result(None)
-
-    timer_handle = loop.call_later(timeout, _handle_timeout)
-    condition_wait = loop.create_task(condition.wait())
-
-    def _handle_wait_complete(_: asyncio.Task) -> None:
-        if not future.done():
-            future.set_result(None)
-
-    condition_wait.add_done_callback(_handle_wait_complete)
-
-    try:
-        await future
-    finally:
-        timer_handle.cancel()
-        if not condition_wait.done():
-            condition_wait.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await condition_wait
-
-
 class _AsyncSender(threading.Thread):
     """A thread to handle sending DNSOutgoing for asyncio."""
 

diff --git a/zeroconf/utils/__init__.py b/zeroconf/utils/__init__.py
@@ -0,0 +1,21 @@
+""" Multicast DNS Service Discovery for Python, v0.14-wmcbrine
+    Copyright 2003 Paul Scott-Murphy, 2014 William McBrine
+
+    This module provides a framework for the use of DNS Service Discovery
+    using IP multicast.
+
+    This library is free software; you can redistribute it and/or
+    modify it under the terms of the GNU Lesser General Public
+    License as published by the Free Software Foundation; either
+    version 2.1 of the License, or (at your option) any later version.
+
+    This library is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+    Lesser General Public License for more details.
+
+    You should have received a copy of the GNU Lesser General Public
+    License along with this library; if not, write to the Free Software
+    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301
+    USA
+"""
diff --git a/zeroconf/utils/aio.py b/zeroconf/utils/aio.py
@@ -0,0 +1,67 @@
+""" Multicast DNS Service Discovery for Python, v0.14-wmcbrine
+    Copyright 2003 Paul Scott-Murphy, 2014 William McBrine
+
+    This module provides a framework for the use of DNS Service Discovery
+    using IP multicast.
+
+    This library is free software; you can redistribute it and/or
+    modify it under the terms of the GNU Lesser General Public
+    License as published by the Free Software Foundation; either
+    version 2.1 of the License, or (at your option) any later version.
+
+    This library is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+    Lesser General Public License for more details.
+
+    You should have received a copy of the GNU Lesser General Public
+    License along with this library; if not, write to the Free Software
+    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301
+    USA
+"""
+
+import asyncio
+import contextlib
+from typing import Optional, cast
+
+
+# Switch to asyncio.wait_for once https://bugs.python.org/issue39032 is fixed
+async def wait_condition_or_timeout(condition: asyncio.Condition, timeout: float) -> None:
+    """Wait for a condition or timeout."""
+    loop = asyncio.get_event_loop()
+    future = loop.create_future()
+
+    def _handle_timeout() -> None:
+        if not future.done():
+            future.set_result(None)
+
+    timer_handle = loop.call_later(timeout, _handle_timeout)
+    condition_wait = loop.create_task(condition.wait())
+
+    def _handle_wait_complete(_: asyncio.Task) -> None:
+        if not future.done():
+            future.set_result(None)
+
+    condition_wait.add_done_callback(_handle_wait_complete)
+
+    try:
+        await future
+    finally:
+        timer_handle.cancel()
+        if not condition_wait.done():
+            condition_wait.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await condition_wait
+
+
+# Remove the call to _get_running_loop once we drop python 3.6 support
+def get_running_loop() -> Optional[asyncio.AbstractEventLoop]:
+    """Check if an event loop is already running."""
+    with contextlib.suppress(RuntimeError):
+        if hasattr(asyncio, "get_running_loop"):
+            return cast(
+                asyncio.AbstractEventLoop,
+                asyncio.get_running_loop(),  # type: ignore  # pylint: disable=no-member  # noqa
+            )
+        return asyncio._get_running_loop()  # pylint: disable=no-member,protected-access
+    return None