path: root/tools/write_pyi.py

from argparse import ArgumentParser
from pathlib import Path
import re
import sys
from tempfile import NamedTemporaryFile
import textwrap
import typing

from mako.pygen import PythonPrinter

sys.path.append(str(Path(__file__).parent.parent))

if True:  # avoid flake/zimports messing with the order
    from alembic.operations.base import Operations
    from alembic.runtime.environment import EnvironmentContext
    from alembic.runtime.migration import MigrationContext
    from alembic.script.write_hooks import console_scripts
    from alembic.util.compat import inspect_formatargspec
    from alembic.util.compat import inspect_getfullargspec
    from alembic.operations import ops
    import sqlalchemy as sa

IGNORE_ITEMS = {
    "op": {"context", "create_module_class_proxy"},
    "context": {
        "create_module_class_proxy",
        "get_impl",
        "requires_connection",
    },
}
TRIM_MODULE = [
    "alembic.runtime.migration.",
    "alembic.operations.base.",
    "alembic.operations.ops.",
    "sqlalchemy.engine.base.",
    "sqlalchemy.engine.url.",
    "sqlalchemy.sql.schema.",
    "sqlalchemy.sql.selectable.",
    "sqlalchemy.sql.elements.",
    "sqlalchemy.sql.type_api.",
    "sqlalchemy.sql.functions.",
    "sqlalchemy.sql.dml.",
]
CONTEXT_MANAGERS = {"op": ["batch_alter_table"]}
ADDITIONAL_ENV = {"MigrationContext": MigrationContext}


def generate_pyi_for_proxy(
    cls: type,
    progname: str,
    source_path: Path,
    destination_path: Path,
    ignore_output: bool,
    file_key: str,
):
    ignore_items = IGNORE_ITEMS.get(file_key, set())
    context_managers = CONTEXT_MANAGERS.get(file_key, [])
    if sys.version_info < (3, 11):
        raise RuntimeError(
            "This script must be run with Python 3.11 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
    with open(source_path) as read_file:
        for line in read_file:
            if line.startswith("# ### this file stubs are generated by"):
                read_imports = True
            elif line.startswith("### end imports ###"):
                read_imports = False
                break
            elif read_imports:
                imports.append(line.rstrip())

    with open(destination_path, "w") as buf:
        printer = PythonPrinter(buf)

        printer.writeline(
            f"# ### this file stubs are generated by {progname} "
            "- do not edit ###"
        )
        for line in imports:
            buf.write(line + "\n")
        printer.writeline("### end imports ###")
        buf.write("\n\n")

        module = sys.modules[cls.__module__]
        env = {
            **typing.__dict__,
            **sa.sql.schema.__dict__,
            **sa.__dict__,
            **sa.types.__dict__,
            **ADDITIONAL_ENV,
            **ops.__dict__,
            **module.__dict__,
        }

        for name in dir(cls):
            if name.startswith("_") or name in ignore_items:
                continue
            meth = getattr(cls, name, None)
            if callable(meth):
                # If there are overloads, generate only those
                # Do not generate the base implementation to avoid mypy errors
                overloads = typing.get_overloads(meth)
                if overloads:
                    # use enumerate so we can generate docs on the last overload
                    for i, ovl in enumerate(overloads, 1):
                        _generate_stub_for_meth(
                            ovl,
                            cls,
                            printer,
                            env,
                            is_context_manager=name in context_managers,
                            is_overload=True,
                            base_method=meth,
                            gen_docs=(i == len(overloads)),
                        )
                else:
                    _generate_stub_for_meth(
                        meth,
                        cls,
                        printer,
                        env,
                        is_context_manager=name in context_managers,
                    )
            else:
                _generate_stub_for_attr(cls, name, printer, env)

        printer.close()

    console_scripts(
        str(destination_path),
        {"entrypoint": "zimports", "options": "-e"},
        ignore_output=ignore_output,
    )
    # note that we do not distribute pyproject.toml with the distribution
    # right now due to user complaints, so we can't refer to it here because
    # this all has to run as part of the test suite
    console_scripts(
        str(destination_path),
        {"entrypoint": "black", "options": "-l79"},
        ignore_output=ignore_output,
    )


def _generate_stub_for_attr(cls, name, printer, env):
    try:
        annotations = typing.get_type_hints(cls, env)
    except NameError:
        annotations = cls.__annotations__
    type_ = annotations.get(name, "Any")
    if isinstance(type_, str) and type_[0] in "'\"":
        type_ = type_[1:-1]
    printer.writeline(f"{name}: {type_}")


def _generate_stub_for_meth(
    fn,
    cls,
    printer,
    env,
    is_context_manager,
    is_overload=False,
    base_method=None,
    gen_docs=True,
):
    while hasattr(fn, "__wrapped__"):
        fn = fn.__wrapped__

    name = fn.__name__
    spec = inspect_getfullargspec(fn)
    try:
        annotations = typing.get_type_hints(fn, env)
        spec.annotations.update(annotations)
    except NameError as e:
        print(f"{cls.__name__}.{name} NameError: {e}", file=sys.stderr)

    name_args = spec[0]
    assert name_args[0:1] == ["self"] or name_args[0:1] == ["cls"]

    name_args[0:1] = []

    def _formatannotation(annotation, base_module=None):
        if getattr(annotation, "__module__", None) == "typing":
            retval = repr(annotation).replace("typing.", "")
        elif isinstance(annotation, type):
            retval = annotation.__qualname__
        else:
            retval = annotation

        for trim in TRIM_MODULE:
            retval = retval.replace(trim, "")

        retval = re.sub(
            r'ForwardRef\(([\'"].+?[\'"])\)', lambda m: m.group(1), retval
        )
        retval = re.sub("NoneType", "None", retval)
        return retval

    def _formatvalue(value):
        return "=" + ("..." if value is Ellipsis else repr(value))

    argspec = inspect_formatargspec(
        *spec,
        formatannotation=_formatannotation,
        formatvalue=_formatvalue,
        formatreturns=lambda val: f"-> {_formatannotation(val)}",
    )

    overload = "@overload" if is_overload else ""
    contextmanager = "@contextmanager" if is_context_manager else ""

    fn_doc = base_method.__doc__ if base_method else fn.__doc__
    has_docs = gen_docs and fn_doc is not None
    string_prefix = "r" if chr(92) in fn_doc else ""
    docs = f'{string_prefix}"""' + f"{fn_doc}" + '"""' if has_docs else ""

    func_text = textwrap.dedent(
        f"""
    {overload}
    {contextmanager}
    def {name}{argspec}: {"..." if not docs else ""}
        {docs}
    """
    )

    printer.write_indented_block(func_text)


def run_file(
    source_path: Path, cls_to_generate: type, stdout: bool, file_key: str
):
    progname = Path(sys.argv[0]).as_posix()
    if not stdout:
        generate_pyi_for_proxy(
            cls_to_generate,
            progname,
            source_path=source_path,
            destination_path=source_path,
            ignore_output=False,
            file_key=file_key,
        )
    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_path,
                ignore_output=True,
                file_key=file_key,
            )
            sys.stdout.write(f_path.read_text())
        f_path.unlink()


def main(args):
    location = Path(__file__).parent.parent / "alembic"
    if args.file in {"all", "op"}:
        run_file(location / "op.pyi", Operations, args.stdout, "op")
    if args.file in {"all", "context"}:
        run_file(
            location / "context.pyi",
            EnvironmentContext,
            args.stdout,
            "context",
        )


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--file",
        choices={"op", "context", "all"},
        default="all",
        help="Which file to generate. Default is to regenerate all files",
    )
    parser.add_argument(
        "--stdout",
        action="store_true",
        help="Write to stdout instead of saving to file",
    )
    args = parser.parse_args()
    main(args)