Merge "Implement rudimentary asyncio support w/ asyncpg"

2 files changed, 61 insertions, 2 deletions

diff --git a/lib/sqlalchemy/testing/plugin/plugin_base.py b/lib/sqlalchemy/testing/plugin/plugin_base.py
index b31a4ff3e..49ff0f975 100644
--- a/lib/sqlalchemy/testing/plugin/plugin_base.py
+++ b/lib/sqlalchemy/testing/plugin/plugin_base.py
@@ -48,7 +48,6 @@ testing = None
 util = None
 file_config = None
 
-
 logging = None
 include_tags = set()
 exclude_tags = set()
@@ -193,6 +192,12 @@ def setup_options(make_option):
         default=False,
         help="Unconditionally write/update profiling data.",
     )
+    make_option(
+        "--dump-pyannotate",
+        type=str,
+        dest="dump_pyannotate",
+        help="Run pyannotate and dump json info to given file",
+    )
 
 
 def configure_follower(follower_ident):
@@ -378,7 +383,6 @@ def _engine_uri(options, file_config):
         cfg = provision.setup_config(
             db_url, options, file_config, provision.FOLLOWER_IDENT
         )
-
         if not config._current:
             cfg.set_as_current(cfg, testing)
 
diff --git a/lib/sqlalchemy/testing/plugin/pytestplugin.py b/lib/sqlalchemy/testing/plugin/pytestplugin.py
index 015598952..3df239afa 100644
--- a/lib/sqlalchemy/testing/plugin/pytestplugin.py
+++ b/lib/sqlalchemy/testing/plugin/pytestplugin.py
@@ -26,6 +26,11 @@ else:
         from typing import Sequence
 
 try:
+    import asyncio
+except ImportError:
+    pass
+
+try:
     import xdist  # noqa
 
     has_xdist = True
@@ -101,6 +106,24 @@ def pytest_configure(config):
 
     plugin_base.set_fixture_functions(PytestFixtureFunctions)
 
+    if config.option.dump_pyannotate:
+        global DUMP_PYANNOTATE
+        DUMP_PYANNOTATE = True
+
+
+DUMP_PYANNOTATE = False
+
+
+@pytest.fixture(autouse=True)
+def collect_types_fixture():
+    if DUMP_PYANNOTATE:
+        from pyannotate_runtime import collect_types
+
+        collect_types.start()
+    yield
+    if DUMP_PYANNOTATE:
+        collect_types.stop()
+
 
 def pytest_sessionstart(session):
     plugin_base.post_begin()
@@ -109,6 +132,31 @@ def pytest_sessionstart(session):
 def pytest_sessionfinish(session):
     plugin_base.final_process_cleanup()
 
+    if session.config.option.dump_pyannotate:
+        from pyannotate_runtime import collect_types
+
+        collect_types.dump_stats(session.config.option.dump_pyannotate)
+
+
+def pytest_collection_finish(session):
+    if session.config.option.dump_pyannotate:
+        from pyannotate_runtime import collect_types
+
+        lib_sqlalchemy = os.path.abspath("lib/sqlalchemy")
+
+        def _filter(filename):
+            filename = os.path.normpath(os.path.abspath(filename))
+            if "lib/sqlalchemy" not in os.path.commonpath(
+                [filename, lib_sqlalchemy]
+            ):
+                return None
+            if "testing" in filename:
+                return None
+
+            return filename
+
+        collect_types.init_types_collection(filter_filename=_filter)
+
 
 if has_xdist:
     import uuid
@@ -518,3 +566,10 @@ class PytestFixtureFunctions(plugin_base.FixtureFunctions):
 
     def get_current_test_name(self):
         return os.environ.get("PYTEST_CURRENT_TEST")
+
+    def async_test(self, fn):
+        @_pytest_fn_decorator
+        def decorate(fn, *args, **kwargs):
+            asyncio.get_event_loop().run_until_complete(fn(*args, **kwargs))
+
+        return decorate(fn)