clean up compat module

hyperglass/compat/__init__.py

@@ -1 +1,10 @@
-"""Functions for maintaining compatability with older Python versions or libraries."""
+"""Functions for maintaining compatibility with older Python versions or libraries."""
+
+# Local
+from ._sshtunnel import SSHTunnelForwarder, BaseSSHTunnelForwarderError, open_tunnel
+
+__all__ = (
+    "BaseSSHTunnelForwarderError",
+    "open_tunnel",
+    "SSHTunnelForwarder",
+)

hyperglass/compat/_asyncio.py

@@ -1,116 +0,0 @@
-"""Functions for maintaining asyncio compatability with other versions of Python."""
-
-# Standard Library
-import sys
-import asyncio
-import weakref
-
-try:
-    # Standard Library
-    from asyncio import get_running_loop
-except ImportError:
-    # Standard Library
-    from asyncio.events import _get_running_loop as get_running_loop
-
-RUNNING_PYTHON_VERSION = sys.version_info
-
-# _patch_loop, _patched_run, and _cancel_all_tasks are taken directly
-# from github.com/nickdavis:
-# https://gist.github.com/nickdavies/4a37c6cd9dcc7041fddd2d2a81cee383
-
-# These functions are a backport of the functionality added in
-# Python 3.7 to support asyncio.run(), which is used in several areas
-# of hyperglass. Because the LTS version of Ubuntu at this time (18.04)
-# still ships with Python 3.6, compatibility with Python 3.6 is the
-# goal.
-
-
-def _patch_loop(loop):
-    tasks = weakref.WeakSet()
-
-    task_factory = [None]
-
-    def _set_task_factory(factory):
-        task_factory[0] = factory
-
-    def _get_task_factory():
-        return task_factory[0]
-
-    def _safe_task_factory(loop, coro):
-        if task_factory[0] is None:
-            task = asyncio.Task(coro, loop=loop)
-            if task._source_traceback:
-                del task._source_traceback[-1]
-        else:
-            task = task_factory[0](loop, coro)
-        tasks.add(task)
-        return task
-
-    loop.set_task_factory(_safe_task_factory)
-    loop.set_task_factory = _set_task_factory
-    loop.get_task_factory = _get_task_factory
-
-    return tasks
-
-
-def _cancel_all_tasks(loop, tasks):
-    to_cancel = [task for task in tasks if not task.done()]
-
-    if not to_cancel:
-        return
-
-    for task in to_cancel:
-        task.cancel()
-
-    loop.run_until_complete(asyncio.gather(*to_cancel, loop=loop, return_exceptions=True))
-
-    for task in to_cancel:
-        if task.cancelled():
-            continue
-        if task.exception() is not None:
-            loop.call_exception_handler(
-                {
-                    "message": "unhandled exception during asyncio.run() shutdown",
-                    "exception": task.exception(),
-                    "task": task,
-                }
-            )
-
-
-def _patched_run(main, *, debug=False):
-    try:
-        loop = get_running_loop()
-    except RuntimeError:
-        loop = None
-
-    if loop is not None:
-        raise RuntimeError("asyncio.run() cannot be called from a running event loop")
-
-    if not asyncio.iscoroutine(main):
-        raise ValueError("a coroutine was expected, got {!r}".format(main))
-
-    loop = asyncio.new_event_loop()
-    tasks = _patch_loop(loop)
-
-    try:
-        asyncio.set_event_loop(loop)
-        loop.set_debug(debug)
-        return loop.run_until_complete(main)
-    finally:
-        try:
-            _cancel_all_tasks(loop, tasks)
-            loop.run_until_complete(loop.shutdown_asyncgens())
-        finally:
-            asyncio.set_event_loop(None)
-            loop.close()
-
-
-# If local system's python version is at least 3.6, use the backported
-# asyncio runner.
-if RUNNING_PYTHON_VERSION >= (3, 6):
-    aiorun = _patched_run
-
-# If the local system's python version is at least 3.7, use the standard
-# library's asyncio.run()
-elif RUNNING_PYTHON_VERSION >= (3, 7):
-    aiorun = asyncio.run

hyperglass/execution/drivers/_common.py

@@ -13,9 +13,9 @@ from ._construct import Construct
 
 if t.TYPE_CHECKING:
     # Project
+    from hyperglass.compat import SSHTunnelForwarder
     from hyperglass.models.api import Query
     from hyperglass.models.data import OutputDataModel
-    from hyperglass.compat._sshtunnel import SSHTunnelForwarder
     from hyperglass.models.config.devices import Device
 
 

hyperglass/execution/drivers/ssh.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
 # Project
 from hyperglass.log import log
 from hyperglass.state import use_state
-from hyperglass.compat._sshtunnel import BaseSSHTunnelForwarderError, open_tunnel
+from hyperglass.compat import BaseSSHTunnelForwarderError, open_tunnel
 from hyperglass.exceptions.public import ScrapeError
 
 # Local
@@ -14,7 +14,7 @@ from ._common import Connection
 
 if TYPE_CHECKING:
     # Project
-    from hyperglass.compat._sshtunnel import SSHTunnelForwarder
+    from hyperglass.compat import SSHTunnelForwarder
 
 
 class SSHConnection(Connection):