"""The pyi generator module."""
import importlib
import inspect
import os
import re
import sys
from inspect import getfullargspec
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, get_args # NOQA
import black
from reflex.components.component import Component
from reflex.vars import Var
ruff_dont_remove = [Var, Optional, Dict, List]
EXCLUDED_FILES = [
"__init__.py",
"component.py",
"bare.py",
"foreach.py",
"cond.py",
"multiselect.py",
]
DEFAULT_TYPING_IMPORTS = {"overload", "Optional", "Union"}
def _get_type_hint(value, top_level=True, no_union=False):
res = ""
args = get_args(value)
if args:
res = f"{value.__name__}[{', '.join([_get_type_hint(arg, top_level=False) for arg in args if arg is not type(None)])}]"
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):
return [
line.split(" = ")[0]
for line in inspect.getsource(_module).splitlines()
if line.startswith(_var_name)
]
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",
]
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(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}, "
for name, value in _class.__annotations__.items():
if name in create_spec.kwonlyargs:
continue
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__}': # type: ignore\n"
definition += self._generate_docstrings(_class, _class.__annotations__.keys())
lines.append(definition)
lines.append(" ...")
return lines
def _generate_docstrings(self, _class, _props):
props_comments = {}
comments = []
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:
# This isn't a prop, so continue.
props_comments[prop] = "\n".join(
[comment.strip().strip("#") for comment in comments]
)
comments.clear()
continue
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.extend(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 = [
(name, obj)
for name, obj in vars(self.current_module).items()
if not name.startswith("__")
and not inspect.isclass(obj)
and not inspect.isfunction(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)