diff options
Diffstat (limited to 'tools/write_pyi.py')
| -rw-r--r-- | tools/write_pyi.py | 18 |
1 files changed, 12 insertions, 6 deletions
diff --git a/tools/write_pyi.py b/tools/write_pyi.py index 61c3052..60728a8 100644 --- a/tools/write_pyi.py +++ b/tools/write_pyi.py @@ -9,7 +9,7 @@ from mako.pygen import PythonPrinter sys.path.append(str(Path(__file__).parent.parent)) -if True: # avoid flake/zimports missing with the order +if True: # avoid flake/zimports messing with the order from alembic.operations.base import Operations from alembic.runtime.environment import EnvironmentContext from alembic.script.write_hooks import console_scripts @@ -19,11 +19,9 @@ if True: # avoid flake/zimports missing with the order IGNORE_ITEMS = { "op": {"context", "create_module_class_proxy"}, "context": { - "config", "create_module_class_proxy", "get_impl", "requires_connection", - "script", }, } @@ -36,6 +34,13 @@ def generate_pyi_for_proxy( ignore_output: bool, ignore_items: set, ): + if sys.version_info < (3, 9): + raise RuntimeError("This script must be run with Python 3.9 or higher") + + # When using an absolute path on windows, this will generate the correct + # relative path that shall be written to the top comment of the pyi file. + if Path(progname).is_absolute(): + progname = Path(progname).relative_to(Path().cwd()).as_posix() imports = [] read_imports = False @@ -88,7 +93,8 @@ def generate_pyi_for_proxy( def _generate_stub_for_attr(cls, name, printer): - printer.writeline(f"{name}: Any") + type_ = cls.__annotations__.get(name, "Any") + printer.writeline(f"{name}: {type_}") def _generate_stub_for_meth(cls, name, printer): @@ -154,15 +160,15 @@ def run_file( else: with NamedTemporaryFile(delete=False, suffix=".pyi") as f: f.close() + f_path = Path(f.name) generate_pyi_for_proxy( cls_to_generate, progname, source_path=source_path, - destination_path=f.name, + destination_path=f_path, ignore_output=True, ignore_items=ignore_items, ) - f_path = Path(f.name) sys.stdout.write(f_path.read_text()) f_path.unlink() |