feat: generate __init__.pyi for mypy

cpcloud · Apr 4, 2022 · 9d7fe5b · 9d7fe5b
1 parent 95f697f
commit 9d7fe5b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/protoletariat/fdsetgen.py b/protoletariat/fdsetgen.py
@@ -95,11 +95,22 @@ def fix_imports(
 
         if create_package:
             python_out.joinpath("__init__.py").touch(exist_ok=True)
+            _create_pyi_init(python_out)
 
             # recursively create packages
             for dir_entry in python_out.rglob("*"):
                 if dir_entry.is_dir():
                     dir_entry.joinpath("__init__.py").touch(exist_ok=True)
+                    _create_pyi_init(dir_entry)
+
+
+def _create_pyi_init(root: Path) -> None:
+    with root.joinpath("__init__.pyi").open(mode="a") as f:
+        f.writelines(
+            f"from . import {path.stem}\n"
+            for path in root.glob("*")
+            if path.stem != "__init__"
+        )
 
 
 class Protoc(FileDescriptorSetGenerator):

diff --git a/protoletariat/tests/test_fix_imports.py b/protoletariat/tests/test_fix_imports.py
@@ -121,6 +121,11 @@ def test_grpc_no_imports(  # type: ignore[misc]
     result = no_imports_service.generate(cli, args=["--in-place", "--create-package"])
     assert result.exit_code == 0
 
+    init = no_imports_service.package_dir.joinpath("__init__.pyi")
+    init_lines = init.read_text().splitlines()
+    assert "from . import no_imports_service_pb2" in init_lines
+    assert "from . import no_imports_service_pb2_grpc" in init_lines
+
     service_module_types = no_imports_service.package_dir.joinpath(
         "no_imports_service_pb2_grpc.pyi"
     )