From ee9c95d56f6db85747d77d59fdded8c13d83560f Mon Sep 17 00:00:00 2001 From: JasonGrace2282 Date: Mon, 1 Apr 2024 16:11:26 -0400 Subject: [PATCH] fix(autoaliasattr): search for type aliases under if TYPE_CHECKING --- manim/utils/docbuild/module_parsing.py | 119 ++++++++++++++++--------- 1 file changed, 77 insertions(+), 42 deletions(-) diff --git a/manim/utils/docbuild/module_parsing.py b/manim/utils/docbuild/module_parsing.py index 46c31d41d0..8ff7c0fe5e 100644 --- a/manim/utils/docbuild/module_parsing.py +++ b/manim/utils/docbuild/module_parsing.py @@ -44,6 +44,10 @@ MANIM_ROOT = Path(__file__).resolve().parent.parent.parent +# In the following, we will use ``type(xyz) is xyz_type`` instead of +# isinstance checks to make sure no subclasses of the type pass the +# check + def parse_module_attributes() -> tuple[AliasDocsDict, DataDict]: """Read all files, generate Abstract Syntax Trees from them, and @@ -70,7 +74,7 @@ def parse_module_attributes() -> tuple[AliasDocsDict, DataDict]: for module_path in MANIM_ROOT.rglob("*.py"): module_name = module_path.resolve().relative_to(MANIM_ROOT) module_name = list(module_name.parts) - module_name[-1] = module_name[-1][:-3] # remove .py + module_name[-1] = module_name[-1].removesuffix(".py") module_name = ".".join(module_name) module_content = module_path.read_text(encoding="utf-8") @@ -107,53 +111,84 @@ def parse_module_attributes() -> tuple[AliasDocsDict, DataDict]: data_list.append(data_name) continue - # If we encounter an assignment annotated as "TypeAlias": + # if it's defined under if TYPE_CHECKING + # go through the body of the if statement if ( - type(node) is ast.AnnAssign - and type(node.annotation) is ast.Name - and node.annotation.id == "TypeAlias" - and type(node.target) is ast.Name - and node.value is not None + # NOTE: This logic does not (and cannot) + # check if the comparison is against a + # variable called TYPE_CHECKING + # It also says that you cannot do the following + # import typing as foo + # if foo.TYPE_CHECKING: + # BAR: TypeAlias = ... + type(node) is ast.If + and ( + ( + # if TYPE_CHECKING + type(node.test) is ast.Name + and node.test.id == "TYPE_CHECKING" + ) + or ( + # if typing.TYPE_CHECKING + type(node.test) is ast.Attribute + and type(node.test.value) is ast.Name + and node.test.value.id == "typing" + and node.test.attr == "TYPE_CHECKING" + ) + ) ): - alias_name = node.target.id - def_node = node.value - # If it's an Union, replace it with vertical bar notation + inner_nodes = node.body + else: + inner_nodes = [node] + + for node in inner_nodes: + # If we encounter an assignment annotated as "TypeAlias": if ( - type(def_node) is ast.Subscript - and type(def_node.value) is ast.Name - and def_node.value.id == "Union" + type(node) is ast.AnnAssign + and type(node.annotation) is ast.Name + and node.annotation.id == "TypeAlias" + and type(node.target) is ast.Name + and node.value is not None ): - definition = " | ".join( - ast.unparse(elem) for elem in def_node.slice.elts - ) + alias_name = node.target.id + def_node = node.value + # If it's an Union, replace it with vertical bar notation + if ( + type(def_node) is ast.Subscript + and type(def_node.value) is ast.Name + and def_node.value.id == "Union" + ): + definition = " | ".join( + ast.unparse(elem) for elem in def_node.slice.elts + ) + else: + definition = ast.unparse(def_node) + + definition = definition.replace("npt.", "") + if category_dict is None: + module_dict[""] = {} + category_dict = module_dict[""] + category_dict[alias_name] = {"definition": definition} + alias_info = category_dict[alias_name] + continue + + # If here, the node is not a TypeAlias definition + alias_info = None + + # It could still be a module attribute definition. + # Does the assignment have a target of type Name? Then + # it could be considered a definition of a module attribute. + if type(node) is ast.AnnAssign: + target = node.target + elif type(node) is ast.Assign and len(node.targets) == 1: + target = node.targets[0] else: - definition = ast.unparse(def_node) - - definition = definition.replace("npt.", "") - if category_dict is None: - module_dict[""] = {} - category_dict = module_dict[""] - category_dict[alias_name] = {"definition": definition} - alias_info = category_dict[alias_name] - continue - - # If here, the node is not a TypeAlias definition - alias_info = None - - # It could still be a module attribute definition. - # Does the assignment have a target of type Name? Then - # it could be considered a definition of a module attribute. - if type(node) is ast.AnnAssign: - target = node.target - elif type(node) is ast.Assign and len(node.targets) == 1: - target = node.targets[0] - else: - target = None + target = None - if type(target) is ast.Name: - data_name = target.id - else: - data_name = None + if type(target) is ast.Name: + data_name = target.id + else: + data_name = None if len(module_dict) > 0: ALIAS_DOCS_DICT[module_name] = module_dict