gen_pyi.py

from __future__ import annotations

import ast
import importlib
import inspect
import pathlib
import re
import sys
import warnings
from collections import OrderedDict
from typing import Callable

# Import registry of decorated classes
from sift_client._internal.sync_wrapper import SyncAPIRegistration, _registered

FUTURE_IMPORTS = "from __future__ import annotations"
TYPE_CHECKING_IMPORT = "from typing import TYPE_CHECKING"
TYPE_CHECK_BLOCK = "if TYPE_CHECKING:"
HEADER = """\
# Auto-generated stub
"""

CLASS_TEMPLATE = """
class {cls_name}:
{doc}
{methods}
"""

METHOD_TEMPLATE = """\
    {decorator}
    def {meth_name}(self{params}){ret_ann}:
{docstring_section}
        ...
"""


def extract_imports(path: pathlib.Path) -> list[str]:
    """Parse the given Python file and return a list of its import statements (as strings)."""
    source = path.read_text()
    tree = ast.parse(source, filename=str(path))

    imports = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                stmt = f"import {alias.name}" + (f" as {alias.asname}" if alias.asname else "")
                imports.append(stmt)
        elif isinstance(node, ast.ImportFrom):
            module = node.module or ""
            level = "." * node.level
            names = ", ".join(
                alias.name + (f" as {alias.asname}" if alias.asname else "") for alias in node.names
            )
            stmt = f"from {level + module} import {names}"
            imports.append(stmt)
    return imports


def format_annotation(ann):
    # Handle multiple ways annotations are represented
    if ann is inspect._empty:
        return ""
    if isinstance(ann, str):
        return ann
    origin = getattr(ann, "__origin__", None)
    if origin:
        args = ", ".join(format_annotation(a) for a in ann.__args__)
        name = getattr(origin, "__name__", repr(origin))
        return f"{name}[{args}]"
    if hasattr(ann, "__name__"):
        return ann.__name__
    return getattr(ann, "__qualname__", repr(ann))


def write_stub_files(stub_files: dict[pathlib.Path, str]):
    for pyi_file, content in stub_files.items():
        print(f"Writing stub: {pyi_file}")
        pyi_file.write_text(content)


# def generate_stub_for_class(classes: SyncAPIRegistration) -> str:


def generate_stubs_for_module(path_arg: str | pathlib.Path) -> dict[pathlib.Path, str]:
    cwd = pathlib.Path.cwd().resolve()
    candidate = pathlib.Path(path_arg)
    abs_path = (cwd / candidate).resolve()

    if abs_path.is_file():
        search_root = abs_path.parent
    elif abs_path.is_dir():
        search_root = abs_path
    else:
        raise ValueError(f"{path_arg!r} is neither a file nor a directory")

    stub_files: dict[pathlib.Path, str] = {}

    # Find all python files in the directory
    for py_file in search_root.rglob("*.py"):
        if py_file.name.startswith("test_"):
            continue
        # Determine module name to import
        rel = py_file.with_suffix("").relative_to(cwd)
        module_name = ".".join(rel.parts)
        module = importlib.import_module(module_name)
        new_module_imports: list[str] = []

        lines = []
        needs_builtins_import = False

        # Process only classes generated by @generate_sync_api
        classes = _registered

        for cls_name, cls in inspect.getmembers(module, inspect.isclass):
            matching_registered_classes: list[SyncAPIRegistration] = list(
                filter(lambda c: c["sync_cls"].__name__ == cls_name, classes)
            )
            if len(matching_registered_classes) < 1:
                continue

            matched: SyncAPIRegistration = matching_registered_classes[0]

            async_class = matched.get("async_cls")
            if async_class is None:
                warnings.warn(
                    f"Could not find async class for {cls_name}. Skipping stub generation.",
                    stacklevel=2,
                )
                continue

            # Read imports from the original async class module
            source_file = inspect.getsourcefile(async_class)
            if source_file is None:
                warnings.warn(
                    f"Could not find source file for {async_class.__name__}. Skipping stub generation.",
                    stacklevel=2,
                )
                continue

            orig_path = pathlib.Path(source_file).resolve()
            imports = extract_imports(orig_path)
            new_module_imports = new_module_imports + imports

            # Class docstring
            raw_doc = inspect.getdoc(cls) or ""
            if raw_doc:
                doc = (
                    '    """\n'
                    + "\n".join(f"    {l.strip()}" for l in raw_doc.splitlines())
                    + '\n    """'
                )
            else:
                doc = "    ..."

            methods = []

            # Check if class has a method named 'list' that would shadow builtins.list
            has_list_method = any(
                name == "list" and inspect.isfunction(member)
                for name, member in inspect.getmembers(cls, inspect.isfunction)
            )
            if has_list_method:
                needs_builtins_import = True

            # Method stub generation
            orig_methods = inspect.getmembers(cls, inspect.isfunction)
            for meth_name, method in orig_methods:
                methods.append(generate_method_stub(meth_name, method, module, "", has_list_method))

            # Property stub generation
            orig_properties = inspect.getmembers(cls, lambda o: isinstance(o, property))
            for prop_name, prop in orig_properties:
                # Getters
                if prop.fget:
                    methods.append(generate_method_stub(prop_name, prop.fget, module, "@property", has_list_method))
                # Setters
                if prop.fset:
                    methods.append(
                        generate_method_stub(prop_name, prop.fset, module, "@property.setter", has_list_method)
                    )
                # Deleters
                if prop.fdel:
                    methods.append(
                        generate_method_stub(prop_name, prop.fdel, module, "@property.deleter", has_list_method)
                    )

            stub = CLASS_TEMPLATE.format(
                cls_name=cls_name,
                doc=doc,
                methods="".join(methods),
            )
            lines.append(stub)

        unique_imports = list(OrderedDict.fromkeys(new_module_imports))
        unique_imports.remove(FUTURE_IMPORTS)  # Future imports can't be in type checking block.
        # Add builtins import if any class has a 'list' method (to avoid shadowing)
        if needs_builtins_import and "import builtins" not in unique_imports:
            unique_imports.append("import builtins")
        # Make import block such that all type hints are used in type checking and not actually required.
        import_block = [FUTURE_IMPORTS, TYPE_CHECKING_IMPORT, TYPE_CHECK_BLOCK] + [
            f"    {import_stmt}" for import_stmt in unique_imports
        ]
        lines = [HEADER, *import_block, *lines]
        pyi_file = py_file.with_suffix(".pyi")

        stub_files[pyi_file] = "\n".join(lines)

    return stub_files


def generate_method_stub(name: str, f: Callable, module, decorator: str = "", has_list_method: bool = False) -> str:
    sig = inspect.signature(f)

    # Parameters
    params = []
    has_keyword_only = False
    positional_only_count = 0

    # First pass to count positional-only parameters
    for param in sig.parameters.values():
        if param.kind == inspect.Parameter.POSITIONAL_ONLY:
            positional_only_count += 1

    # Second pass to generate parameter strings
    param_count = 0
    for param in sig.parameters.values():
        if param.name == "self":
            continue

        default = ""
        if param.default is not inspect._empty:
            default = f" = {param.default!r}"

        # Handle different parameter kinds
        if param.kind == inspect.Parameter.VAR_POSITIONAL:
            # Handle *args
            if param.annotation and param.annotation is not inspect._empty:
                params.append(f", *{param.name}: {param.annotation}")
            else:
                params.append(f", *{param.name}")
        elif param.kind == inspect.Parameter.VAR_KEYWORD:
            # Handle **kwargs
            if param.annotation and param.annotation is not inspect._empty:
                params.append(f", **{param.name}: {param.annotation}")
            else:
                params.append(f", **{param.name}")
        elif param.kind == inspect.Parameter.KEYWORD_ONLY:
            # Handle keyword-only parameters (after *)
            if not has_keyword_only:
                # Add a standalone * if this is the first keyword-only parameter
                params.append(", *")
                has_keyword_only = True

            if param.annotation and param.annotation is not inspect._empty:
                params.append(f", {param.name}: {param.annotation}{default}")
            else:
                params.append(f", {param.name}{default}")
        elif param.kind == inspect.Parameter.POSITIONAL_ONLY:
            # Handle positional-only parameters (before /)
            if param.annotation and param.annotation is not inspect._empty:
                params.append(f", {param.name}: {param.annotation}{default}")
            else:
                params.append(f", {param.name}{default}")

            # Add the / separator after the last positional-only parameter
            param_count += 1
            if param_count == positional_only_count:
                params.append(", /")

        else:
            # Handle normal parameters
            if param.annotation and param.annotation is not inspect._empty:
                params.append(f", {param.name}: {param.annotation}{default}")
            else:
                params.append(f", {param.name}{default}")

    params_txt = "".join(params)

    # Return annotation - replace list[ with builtins.list[ if class has a list method
    if sig.return_annotation is not inspect._empty:
        ret_annotation_str = str(sig.return_annotation)
        # Replace list[ with builtins.list[ to avoid shadowing by method named 'list'
        # Use regex to match word boundary before "list[" to avoid false matches
        if has_list_method and "builtins.list[" not in ret_annotation_str:
            # Replace all occurrences of "list[" with "builtins.list["
            ret_annotation_str = re.sub(r'\blist\[', 'builtins.list[', ret_annotation_str)
        ret_txt = f" -> {ret_annotation_str}"
    else:
        ret_txt = ""

    # Method docstring
    raw_mdoc = inspect.getdoc(f) or ""
    if raw_mdoc and raw_mdoc.strip():
        meth_doc = raw_mdoc.replace('"""', '\\"\\"\\"').replace("\n", "\n        ")
        docstring_section = f'        """\n        {meth_doc}\n        """\n'
    else:
        docstring_section = ""

    return METHOD_TEMPLATE.format(
        decorator=decorator,
        meth_name=name,
        params=params_txt,
        ret_ann=ret_txt,
        docstring_section=docstring_section,
    )


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: gen_pyi.py <path/to/module-or-directory>")
        sys.exit(1)

    stubs = generate_stubs_for_module(sys.argv[1])
    write_stub_files(stubs)