reflex/scripts/pyi_generator.py

"""The pyi generator module."""

import ast
import importlib
import inspect
import os
import re
import sys
from inspect import getfullargspec
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Set, Union, get_args  # NOQA

import black

from reflex.components.component import Component

# NOQA
from reflex.components.graphing.recharts.recharts import (
    LiteralAnimationEasing,
    LiteralAreaType,
    LiteralComposedChartBaseValue,
    LiteralDirection,
    LiteralGridType,
    LiteralIconType,
    LiteralIfOverflow,
    LiteralInterval,
    LiteralLayout,
    LiteralLegendAlign,
    LiteralLineType,
    LiteralOrientationTopBottom,
    LiteralOrientationTopBottomLeftRight,
    LiteralPolarRadiusType,
    LiteralPosition,
    LiteralScale,
    LiteralShape,
    LiteralStackOffset,
    LiteralSyncMethod,
    LiteralVerticalAlign,
)
from reflex.components.libs.chakra import (
    LiteralAlertDialogSize,
    LiteralAvatarSize,
    LiteralChakraDirection,
    LiteralColorScheme,
    LiteralDrawerSize,
    LiteralImageLoading,
    LiteralInputVariant,
    LiteralMenuOption,
    LiteralMenuStrategy,
    LiteralTagSize,
)
from reflex.components.radix.themes.base import (
    LiteralAccentColor,
    LiteralAlign,
    LiteralAppearance,
    LiteralGrayColor,
    LiteralJustify,
    LiteralPanelBackground,
    LiteralRadius,
    LiteralScaling,
    LiteralSize,
    LiteralVariant,
)
from reflex.components.radix.themes.components import (
    LiteralButtonSize,
    LiteralSwitchSize,
)
from reflex.components.radix.themes.layout import (
    LiteralBoolNumber,
    LiteralContainerSize,
    LiteralFlexDirection,
    LiteralFlexDisplay,
    LiteralFlexWrap,
    LiteralGridDisplay,
    LiteralGridFlow,
    LiteralSectionSize,
)
from reflex.components.radix.themes.typography import (
    LiteralLinkUnderline,
    LiteralTextAlign,
    LiteralTextSize,
    LiteralTextTrim,
    LiteralTextWeight,
)

# NOQA
from reflex.event import EventChain
from reflex.style import Style
from reflex.utils import format
from reflex.utils import types as rx_types
from reflex.vars import Var

ruff_dont_remove = [
    Var,
    Optional,
    Dict,
    List,
    EventChain,
    Style,
    LiteralInputVariant,
    LiteralColorScheme,
    LiteralChakraDirection,
    LiteralTagSize,
    LiteralDrawerSize,
    LiteralMenuStrategy,
    LiteralMenuOption,
    LiteralAlertDialogSize,
    LiteralAvatarSize,
    LiteralImageLoading,
    LiteralLayout,
    LiteralAnimationEasing,
    LiteralGridType,
    LiteralPolarRadiusType,
    LiteralScale,
    LiteralSyncMethod,
    LiteralStackOffset,
    LiteralComposedChartBaseValue,
    LiteralOrientationTopBottom,
    LiteralAreaType,
    LiteralShape,
    LiteralLineType,
    LiteralDirection,
    LiteralIfOverflow,
    LiteralOrientationTopBottomLeftRight,
    LiteralInterval,
    LiteralLegendAlign,
    LiteralVerticalAlign,
    LiteralIconType,
    LiteralPosition,
    LiteralAccentColor,
    LiteralAlign,
    LiteralAppearance,
    LiteralBoolNumber,
    LiteralButtonSize,
    LiteralContainerSize,
    LiteralFlexDirection,
    LiteralFlexDisplay,
    LiteralFlexWrap,
    LiteralGrayColor,
    LiteralGridDisplay,
    LiteralGridFlow,
    LiteralJustify,
    LiteralLinkUnderline,
    LiteralPanelBackground,
    LiteralRadius,
    LiteralScaling,
    LiteralSectionSize,
    LiteralSize,
    LiteralSwitchSize,
    LiteralTextAlign,
    LiteralTextSize,
    LiteralTextTrim,
    LiteralTextWeight,
    LiteralVariant,
]

EXCLUDED_FILES = [
    "__init__.py",
    "component.py",
    "bare.py",
    "foreach.py",
    "cond.py",
    "multiselect.py",
]

# These props exist on the base component, but should not be exposed in create methods.
EXCLUDED_PROPS = [
    "alias",
    "children",
    "event_triggers",
    "invalid_children",
    "library",
    "lib_dependencies",
    "tag",
    "is_default",
    "special_props",
    "valid_children",
]

DEFAULT_TYPING_IMPORTS = {"overload", "Any", "Dict", "List", "Optional", "Union"}


def _get_type_hint(value, top_level=True, no_union=False):
    res = ""
    args = get_args(value)
    if args:
        inner_container_type_args = (
            [format.wrap(arg, '"') for arg in args]
            if rx_types.is_literal(value)
            else [
                _get_type_hint(arg, top_level=False)
                for arg in args
                if arg is not type(None)
            ]
        )
        res = f"{value.__name__}[{', '.join(inner_container_type_args)}]"

        if value.__name__ == "Var":
            types = [res] + [
                _get_type_hint(arg, top_level=False)
                for arg in args
                if arg is not type(None)
            ]
            if len(types) > 1 and not no_union:
                res = ", ".join(types)
                res = f"Union[{res}]"
    elif isinstance(value, str):
        ev = eval(value)
        res = _get_type_hint(ev, top_level=False) if ev.__name__ == "Var" else value
    else:
        res = value.__name__
    if top_level and not res.startswith("Optional"):
        res = f"Optional[{res}]"
    return res


def _get_typing_import(_module):
    src = [
        line
        for line in inspect.getsource(_module).split("\n")
        if line.startswith("from typing")
    ]
    if len(src):
        return set(src[0].rpartition("from typing import ")[-1].split(", "))
    return set()


def _get_var_definition(_module, _var_name):
    for node in ast.parse(inspect.getsource(_module)).body:
        if isinstance(node, ast.Assign) and _var_name in [
            t.id for t in node.targets if isinstance(t, ast.Name)
        ]:
            return ast.unparse(node)
    raise Exception(f"Could not find var {_var_name} in module {_module}")


class PyiGenerator:
    """A .pyi file generator that will scan all defined Component in Reflex and
    generate the approriate stub.
    """

    modules: list = []
    root: str = ""
    current_module: Any = {}
    default_typing_imports: set = DEFAULT_TYPING_IMPORTS

    def _generate_imports(self, variables, classes):
        variables_imports = {
            type(_var) for _, _var in variables if isinstance(_var, Component)
        }
        bases = {
            base
            for _, _class in classes
            for base in _class.__bases__
            if inspect.getmodule(base) != self.current_module
        } | variables_imports
        bases.add(Component)
        typing_imports = self.default_typing_imports | _get_typing_import(
            self.current_module
        )
        bases = sorted(bases, key=lambda base: base.__name__)
        return [
            f"from typing import {','.join(sorted(typing_imports))}",
            *[f"from {base.__module__} import {base.__name__}" for base in bases],
            "from reflex.vars import Var, BaseVar, ComputedVar",
            "from reflex.event import EventHandler, EventChain, EventSpec",
            "from reflex.style import Style",
        ]

    def _generate_pyi_class(self, _class: type[Component]):
        create_spec = getfullargspec(_class.create)
        lines = [
            "",
            f"class {_class.__name__}({', '.join([base.__name__ for base in _class.__bases__])}):",
        ]
        definition = f"    @overload\n    @classmethod\n    def create(  # type: ignore\n        cls, *children, "

        for kwarg in create_spec.kwonlyargs:
            if kwarg in create_spec.annotations:
                definition += f"{kwarg}: {_get_type_hint(create_spec.annotations[kwarg])} = None, "
            else:
                definition += f"{kwarg}, "

        all_classes = [c for c in _class.__mro__ if issubclass(c, Component)]
        all_props = []
        for target_class in all_classes:
            for name, value in target_class.__annotations__.items():
                if (
                    name in create_spec.kwonlyargs
                    or name in EXCLUDED_PROPS
                    or name in all_props
                ):
                    continue
                all_props.append(name)
                definition += f"{name}: {_get_type_hint(value)} = None, "

        for trigger in sorted(_class().get_event_triggers().keys()):
            definition += f"{trigger}: Optional[Union[EventHandler, EventSpec, List, function, BaseVar]] = None, "

        definition = definition.rstrip(", ")
        definition += f", **props) -> '{_class.__name__}':\n"

        definition += self._generate_docstrings(all_classes, all_props)
        lines.append(definition)
        lines.append("        ...")
        return lines

    def _generate_docstrings(self, _classes, _props):
        props_comments = {}
        comments = []
        for _class in _classes:
            for _i, line in enumerate(inspect.getsource(_class).splitlines()):
                reached_functions = re.search("def ", line)
                if reached_functions:
                    # We've reached the functions, so stop.
                    break

                # Get comments for prop
                if line.strip().startswith("#"):
                    comments.append(line)
                    continue

                # Check if this line has a prop.
                match = re.search("\\w+:", line)
                if match is None:
                    # This line doesn't have a var, so continue.
                    continue

                # Get the prop.
                prop = match.group(0).strip(":")
                if prop in _props:
                    if not comments:  # do not include undocumented props
                        continue
                    props_comments[prop] = "\n".join(
                        [comment.strip().strip("#") for comment in comments]
                    )
                    comments.clear()
                    continue
                if prop in EXCLUDED_PROPS:
                    comments.clear()  # throw away comments for excluded props
        _class = _classes[0]
        new_docstring = []
        for i, line in enumerate(_class.create.__doc__.splitlines()):
            if i == 0:
                new_docstring.append(" " * 8 + '"""' + line)
            else:
                new_docstring.append(line)
            if "*children" in line:
                for nline in [
                    f"{line.split('*')[0]}{n}:{c}" for n, c in props_comments.items()
                ]:
                    new_docstring.append(nline)
        new_docstring += ['"""']
        return "\n".join(new_docstring)

    def _generate_pyi_variable(self, _name, _var):
        return _get_var_definition(self.current_module, _name)

    def _generate_function(self, _name, _func):
        import textwrap

        # Don't generate indented functions.
        source = inspect.getsource(_func)
        if textwrap.dedent(source) != source:
            return []

        definition = "".join([line for line in source.split(":\n")[0].split("\n")])
        return [f"{definition}:", "    ..."]

    def _write_pyi_file(self, variables, functions, classes):
        pyi_content = [
            f'"""Stub file for {self.current_module_path}.py"""',
            "# ------------------- DO NOT EDIT ----------------------",
            "# This file was generated by `scripts/pyi_generator.py`!",
            "# ------------------------------------------------------",
            "",
        ]
        pyi_content.extend(self._generate_imports(variables, classes))

        for _name, _var in variables:
            pyi_content.append(self._generate_pyi_variable(_name, _var))

        for _fname, _func in functions:
            pyi_content.extend(self._generate_function(_fname, _func))

        for _, _class in classes:
            pyi_content.extend(self._generate_pyi_class(_class))

        pyi_filename = f"{self.current_module_path}.pyi"
        pyi_path = os.path.join(self.root, pyi_filename)

        with open(pyi_path, "w") as pyi_file:
            pyi_file.write("\n".join(pyi_content))
        black.format_file_in_place(
            src=Path(pyi_path),
            fast=True,
            mode=black.FileMode(),
            write_back=black.WriteBack.YES,
        )

    def _scan_file(self, file):
        self.current_module_path = os.path.splitext(file)[0]
        module_import = os.path.splitext(os.path.join(self.root, file))[0].replace(
            "/", "."
        )

        self.current_module = importlib.import_module(module_import)

        local_variables = []
        for node in ast.parse(inspect.getsource(self.current_module)).body:
            if isinstance(node, ast.Assign):
                for t in node.targets:
                    if not isinstance(t, ast.Name):
                        # Skip non-var assignment statements
                        continue
                    if t.id.startswith("_"):
                        # Skip private vars
                        continue
                    obj = getattr(self.current_module, t.id, None)
                    if inspect.isclass(obj) or inspect.isfunction(obj):
                        continue
                    local_variables.append((t.id, obj))

        functions = [
            (name, obj)
            for name, obj in vars(self.current_module).items()
            if not name.startswith("__")
            and (
                not inspect.getmodule(obj)
                or inspect.getmodule(obj) == self.current_module
            )
            and inspect.isfunction(obj)
        ]

        class_names = [
            (name, obj)
            for name, obj in vars(self.current_module).items()
            if inspect.isclass(obj)
            and issubclass(obj, Component)
            and obj != Component
            and inspect.getmodule(obj) == self.current_module
        ]
        if not class_names:
            return
        print(f"Parsed {file}: Found {[n for n, _ in class_names]}")
        self._write_pyi_file(local_variables, functions, class_names)

    def _scan_folder(self, folder):
        for root, _, files in os.walk(folder):
            self.root = root
            for file in files:
                if file in EXCLUDED_FILES:
                    continue
                if file.endswith(".py"):
                    self._scan_file(file)

    def scan_all(self, targets):
        """Scan all targets for class inheriting Component and generate the .pyi files.

        Args:
            targets: the list of file/folders to scan.
        """
        for target in targets:
            if target.endswith(".py"):
                self.root, _, file = target.rpartition("/")
                self._scan_file(file)
            else:
                self._scan_folder(target)


if __name__ == "__main__":
    targets = sys.argv[1:] if len(sys.argv) > 1 else ["reflex/components"]
    print(f"Running .pyi generator for {targets}")
    gen = PyiGenerator()
    gen.scan_all(targets)