Add command to generate gino.ext.* stubs (#674)

* run `python -m gino.ext stub` to generate stubs
* fixes #673

Author: fantix

.gitignore

@@ -69,4 +69,7 @@ target/
 .DS_Store
 
 # VS Code
-.vscode
+.vscode
+
+# extension stub files
+src/gino/ext/*.pyi

poetry.lock

@@ -56,6 +56,7 @@ pytest-asyncio = "^0.10.0"
 pytest-mock = "^3.0.0"
 pytest-cov = "^2.8.1"
 black = { version = "^19.10b0", python = ">=3.6" }
+mypy = "^0.770"
 
 # docs
 sphinx = "^3.0.3"

pyproject.toml

@@ -0,0 +1,64 @@
+"""Generate typing stubs for extensions.
+
+    $ python -m gino.ext
+
+"""
+import sys
+import os
+
+try:
+    from importlib.metadata import entry_points
+except ImportError:
+    from importlib_metadata import entry_points
+
+
+if __name__ == "__main__":
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    cmd = sys.argv[1] if len(sys.argv) == 2 else ""
+    eps = list(entry_points().get("gino.extensions", []))
+
+    if cmd == "stub":
+        added = False
+        for ep in eps:
+            path = os.path.join(base_dir, ep.name + ".pyi")
+            if not os.path.exists(path):
+                added = True
+                print("Adding " + path)
+                with open(path, "w") as f:
+                    f.write("from " + ep.value + " import *")
+        if not added:
+            print("Stub files are up to date.")
+
+    elif cmd == "clean":
+        removed = False
+        for filename in os.listdir(base_dir):
+            if filename.endswith(".pyi"):
+                removed = True
+                path = os.path.join(base_dir, filename)
+                print("Removing " + path)
+                os.remove(path)
+        if not removed:
+            print("No stub files found.")
+
+    elif cmd == "list":
+        name_size = max(len(ep.name) for ep in eps)
+        value_size = max(len(ep.value) for ep in eps)
+        for ep in eps:
+            path = os.path.join(base_dir, ep.name + ".pyi")
+            if not os.path.exists(path):
+                path = "no stub file"
+            print(
+                "%s -> gino.ext.%s (%s)"
+                % (ep.value.ljust(value_size), ep.name.ljust(name_size), path)
+            )
+
+    else:
+        print("Manages GINO extensions:")
+        print()
+        print("  python -m gino.ext COMMAND")
+        print()
+        print("Available commands:")
+        print()
+        print("  stub    Generate gino/ext/*.pyi stub files for type checking.")
+        print("  clean   Remove the generated stub files.")
+        print("  list    List installed GINO extensions.")

src/gino/ext/__main__.py

@@ -0,0 +1 @@
+s1 = "111"

src/gino/ext/py.typed

@@ -0,0 +1 @@
+s2 = 222

tests/stub1.py

@@ -2,6 +2,10 @@
 import importlib
 import sys
 import pytest
+import runpy
+from mypy.build import build
+from mypy.modulefinder import BuildSource
+from mypy.options import Options
 
 
 def installed():
@@ -65,3 +69,74 @@ def test_import_error():
     with pytest.raises(ImportError, match="gino-nonexist"):
         # noinspection PyUnresolvedReferences
         from gino.ext import nonexist
+
+
+@pytest.fixture
+def extensions(mocker):
+    EntryPoint = collections.namedtuple("EntryPoint", ["name", "value"])
+    importlib_metadata = mocker.Mock()
+    importlib_metadata.entry_points = lambda: {
+        "gino.extensions": [
+            EntryPoint("demo1", "tests.stub1"),
+            EntryPoint("demo2", "tests.stub2"),
+        ]
+    }
+    mocker.patch.dict("sys.modules", {"importlib.metadata": importlib_metadata})
+
+
+def test_list(mocker, extensions):
+    mocker.patch("sys.argv", ["", "list"])
+    stdout = mocker.patch("sys.stdout.write")
+    runpy.run_module("gino.ext", run_name="__main__")
+    out = "".join(args[0][0] for args in stdout.call_args_list)
+    assert "tests.stub1" in out
+    assert "tests.stub2" in out
+    assert "gino.ext.demo1" in out
+    assert "gino.ext.demo2" in out
+    assert out.count("no stub file") == 2
+
+    mocker.patch("sys.argv", [""])
+    runpy.run_module("gino.ext", run_name="__main__")
+
+
+def test_type_check(mocker, extensions):
+    mocker.patch("sys.argv", ["", "clean"])
+    runpy.run_module("gino.ext", run_name="__main__")
+
+    result = build(
+        [BuildSource(None, None, "from gino.ext.demo3 import s3")], Options()
+    )
+    assert result.errors
+
+    result = build(
+        [BuildSource(None, None, "from gino.ext.demo1 import s1")], Options()
+    )
+    assert result.errors
+
+    mocker.patch("sys.argv", ["", "stub"])
+    runpy.run_module("gino.ext", run_name="__main__")
+    runpy.run_module("gino.ext", run_name="__main__")
+
+    try:
+        result = build(
+            [BuildSource(None, None, "from gino.ext.demo1 import s1")], Options()
+        )
+        assert not result.errors
+
+        result = build(
+            [BuildSource(None, None, "from gino.ext.demo1 import s2")], Options()
+        )
+        assert result.errors
+
+        result = build(
+            [BuildSource(None, None, "from gino.ext.demo2 import s2")], Options()
+        )
+        assert not result.errors
+
+        result = build(
+            [BuildSource(None, None, "from gino.ext.demo2 import s1")], Options()
+        )
+        assert result.errors
+    finally:
+        mocker.patch("sys.argv", ["", "clean"])
+        runpy.run_module("gino.ext", run_name="__main__")