sync_auth0_python.py

import re
import shutil
import sys
from collections.abc import Iterable
from itertools import chain
from pathlib import Path
from subprocess import check_call, run
from textwrap import dedent

ASYNCIFIED_PATH = Path("stubs/auth0-python/auth0/_asyncified")
ASYNCIFY_PYI_PATH = ASYNCIFIED_PATH.parent / "asyncify.pyi"
KEEP_LINES_STARTSWITH = ("from ", "import ", "    def ", "class ", "\n")
AUTO_GENERATED_COMMENT = "# AUTOGENERATED BY scripts/sync_auth0_python.py"

BASE_CLASS_RE = re.compile(r"class (\w+?):")
SUBCLASS_RE = re.compile(r"class (\w+?)\((\w+?)\):")
AS_IMPORT_RE = re.compile(r"(.+?) as \1")
IMPORT_TO_ASYNCIFY_RE = re.compile(r"(from \.\w+? import )(\w+?)\n")
METHOD_TO_ASYNCIFY_RE = re.compile(r"    def (.+?)\(")


def generate_stubs() -> None:
    check_call(("stubgen", "-p", "auth0.authentication", "-p", "auth0.management", "-o", ASYNCIFIED_PATH))

    # Move the generated stubs to the right place
    shutil.copytree(ASYNCIFIED_PATH / "auth0", ASYNCIFIED_PATH, copy_function=shutil.move, dirs_exist_ok=True)
    shutil.rmtree(ASYNCIFIED_PATH / "auth0")

    for path_to_remove in (
        (ASYNCIFIED_PATH / "authentication" / "__init__.pyi"),
        (ASYNCIFIED_PATH / "management" / "__init__.pyi"),
        *chain.from_iterable(
            [
                (already_async, already_async.with_name(already_async.name.removeprefix("async_")))
                for already_async in ASYNCIFIED_PATH.rglob("async_*.pyi")
            ]
        ),
    ):
        path_to_remove.unlink()


def modify_stubs() -> list[tuple[str, str]]:
    base_classes_for_overload: list[tuple[str, str]] = []
    subclasses_for_overload: list[tuple[str, str]] = []

    # Cleanup and modify the stubs
    for stub_path in ASYNCIFIED_PATH.rglob("*.pyi"):
        with stub_path.open() as stub_file:
            lines = stub_file.readlines()
        relative_module = (stub_path.relative_to(ASYNCIFIED_PATH).with_suffix("").as_posix()).replace("/", ".")

        # Only keep imports, classes and public non-special methods
        stub_content = "".join(
            filter(lambda line: "def _" not in line and any(line.startswith(check) for check in KEEP_LINES_STARTSWITH), lines)
        )

        base_classes_for_overload.extend([(relative_module, match) for match in re.findall(BASE_CLASS_RE, stub_content)])
        subclasses_for_overload.extend([(relative_module, groups[0]) for groups in re.findall(SUBCLASS_RE, stub_content)])

        # Remove redundant ` as ` imports
        stub_content = re.sub(AS_IMPORT_RE, "\\1", stub_content)
        # Fix relative imports
        stub_content = stub_content.replace("from ..", "from ...")
        # Rename same-level local imports to use transformed class names ahead of time
        stub_content = re.sub(IMPORT_TO_ASYNCIFY_RE, "\\1_\\2Async\n", stub_content)
        # Prep extra imports
        stub_content = "from typing import type_check_only\n" + stub_content

        # Rename classes to their stub-only asyncified variant and subclass them
        #   Transform subclasses. These are a bit odd since they may have asyncified methods hidden by base class.
        stub_content = re.sub(
            SUBCLASS_RE,
            dedent(
                """\
                @type_check_only
                class _\\1Async(_\\2Async):"""
            ),
            stub_content,
        )
        #   Transform base classes
        stub_content = re.sub(
            BASE_CLASS_RE,
            dedent(
                f"""\
                from auth0.{relative_module} import \\1  # noqa: E402
                @type_check_only
                class _\\1Async(\\1):"""
            ),
            stub_content,
        )
        # Update methods to their asyncified variant
        stub_content = re.sub(METHOD_TO_ASYNCIFY_RE, "    async def \\1_async(", stub_content)
        # Fix empty classes
        stub_content = stub_content.replace("):\n\n", "): ...\n\n")

        stub_path.write_text(f"{AUTO_GENERATED_COMMENT}\n{stub_content}")

    # Broader types last
    return subclasses_for_overload + base_classes_for_overload


def generate_asyncify_pyi(classes_for_overload: Iterable[tuple[str, str]]) -> None:
    imports = ""
    overloads = ""
    for relative_module, class_name in classes_for_overload:
        deduped_class_name = relative_module.replace(".", "_") + class_name
        async_class_name = f"_{class_name}Async"
        deduped_async_class_name = relative_module.replace(".", "_") + async_class_name
        imports += f"from auth0.{relative_module} import {class_name} as {deduped_class_name}\n"
        imports += f"from ._asyncified.{relative_module} import {async_class_name} as {deduped_async_class_name}\n"
        overloads += f"@overload\ndef asyncify(cls: type[{deduped_class_name}]) -> type[{deduped_async_class_name}]: ...\n"

    ASYNCIFY_PYI_PATH.write_text(
        f"""\
{AUTO_GENERATED_COMMENT}
from typing import overload, TypeVar

{imports}
_T = TypeVar("_T")

{overloads}
@overload
def asyncify(cls: type[_T]) -> type[_T]: ...
"""
    )


def main() -> None:
    generate_stubs()
    classes_for_overload = modify_stubs()
    generate_asyncify_pyi(classes_for_overload)

    run((sys.executable, "-m", "pre_commit", "run", "--files", *ASYNCIFIED_PATH.rglob("*.pyi"), ASYNCIFY_PYI_PATH), check=False)


if __name__ == "__main__":
    main()