"""Read and parse all the Manim modules and extract documentation from them.""" from __future__ import annotations import ast import sys from ast import Attribute, Name, Subscript from pathlib import Path from typing import Any, TypeAlias __all__ = ["parse_module_attributes"] AliasInfo: TypeAlias = dict[str, str] """Dictionary with a `definition` key containing the definition of a :class:`TypeAlias` as a string, and optionally a `doc` key containing the documentation for that alias, if it exists. """ AliasCategoryDict: TypeAlias = dict[str, AliasInfo] """Dictionary which holds an `AliasInfo` for every alias name in a same category. """ ModuleLevelAliasDict: TypeAlias = dict[str, AliasCategoryDict] """Dictionary containing every :class:`TypeAlias` defined in a module, classified by category in different `AliasCategoryDict` objects. """ ModuleTypeVarDict: TypeAlias = dict[str, str] """Dictionary containing every :class:`TypeVar` defined in a module.""" AliasDocsDict: TypeAlias = dict[str, ModuleLevelAliasDict] """Dictionary which, for every module in Manim, contains documentation about their module-level attributes which are explicitly defined as :class:`TypeAlias`, separating them from the rest of attributes. """ DataDict: TypeAlias = dict[str, list[str]] """Type for a dictionary which, for every module, contains a list with the names of all their DOCUMENTED module-level attributes (identified by Sphinx via the ``data`` role, hence the name) which are NOT explicitly defined as :class:`TypeAlias`. """ TypeVarDict: TypeAlias = dict[str, ModuleTypeVarDict] """A dictionary mapping module names to dictionaries of :class:`TypeVar` objects.""" ALIAS_DOCS_DICT: AliasDocsDict = {} DATA_DICT: DataDict = {} TYPEVAR_DICT: TypeVarDict = {} 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 # ruff: noqa: E721 def parse_module_attributes() -> tuple[AliasDocsDict, DataDict, TypeVarDict]: """Read all files, generate Abstract Syntax Trees from them, and extract useful information about the type aliases defined in the files: the category they belong to, their definition and their description, separating them from the "regular" module attributes. Returns ------- ALIAS_DOCS_DICT : :class:`AliasDocsDict` A dictionary containing the information from all the type aliases in Manim. See :class:`AliasDocsDict` for more information. DATA_DICT : :class:`DataDict` A dictionary containing the names of all DOCUMENTED module-level attributes which are not a :class:`TypeAlias`. TYPEVAR_DICT : :class:`TypeVarDict` A dictionary containing the definitions of :class:`TypeVar` objects, organized by modules. """ global ALIAS_DOCS_DICT global DATA_DICT global TYPEVAR_DICT if ALIAS_DOCS_DICT or DATA_DICT or TYPEVAR_DICT: return ALIAS_DOCS_DICT, DATA_DICT, TYPEVAR_DICT for module_path in MANIM_ROOT.rglob("*.py"): module_name_t1 = module_path.resolve().relative_to(MANIM_ROOT) module_name_t2 = list(module_name_t1.parts) module_name_t2[-1] = module_name_t2[-1].removesuffix(".py") module_name = ".".join(module_name_t2) module_content = module_path.read_text(encoding="utf-8") # For storing TypeAliases module_dict: ModuleLevelAliasDict = {} category_dict: AliasCategoryDict | None = None alias_info: AliasInfo | None = None # For storing TypeVars module_typevars: ModuleTypeVarDict = {} # For storing regular module attributes data_list: list[str] = [] data_name: str | None = None for node in ast.iter_child_nodes(ast.parse(module_content)): # If we encounter a string: if ( type(node) is ast.Expr and type(node.value) is ast.Constant and type(node.value.value) is str ): string = node.value.value.strip() # It can be the start of a category section_str = "[CATEGORY]" if string.startswith(section_str): category_name = string[len(section_str) :].strip() module_dict[category_name] = {} category_dict = module_dict[category_name] alias_info = None # or a docstring of the alias defined before elif alias_info: alias_info["doc"] = string # or a docstring of the module attribute defined before elif data_name: data_list.append(data_name) continue # if it's defined under if TYPE_CHECKING # go through the body of the if statement if ( # 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" ) ) ): inner_nodes: list[Any] = node.body else: inner_nodes = [node] for node in inner_nodes: # Check if this node is a TypeAlias (type = ) # or an AnnAssign annotated as TypeAlias (: TypeAlias = ). is_type_alias = ( sys.version_info >= (3, 12) and type(node) is ast.TypeAlias ) is_annotated_assignment_with_value = ( 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 ) if is_type_alias or is_annotated_assignment_with_value: # TODO: ast.TypeAlias does not exist before Python 3.12, and that # could be the reason why MyPy does not recognize these as # attributes of node. alias_name = node.name.id if is_type_alias else node.target.id definition_node = node.value # If the definition is a Union, replace with vertical bar notation. # Instead of "Union[Type1, Type2]", we'll have "Type1 | Type2". if ( type(definition_node) is ast.Subscript and type(definition_node.value) is ast.Name and definition_node.value.id == "Union" ): union_elements = definition_node.slice.elts # type: ignore[attr-defined] definition = " | ".join( ast.unparse(elem) for elem in union_elements ) else: definition = ast.unparse(definition_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 # Check if it is a typing.TypeVar ( = TypeVar(...)). elif ( type(node) is ast.Assign and type(node.targets[0]) is ast.Name and type(node.value) is ast.Call and type(node.value.func) is ast.Name and node.value.func.id.endswith("TypeVar") ): module_typevars[node.targets[0].id] = ast.unparse( node.value ).replace("_", r"\_") 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: Name | Attribute | Subscript | ast.expr | None = node.target elif type(node) is ast.Assign and len(node.targets) == 1: target = node.targets[0] else: target = None if type(target) is ast.Name and not ( type(node) is ast.Assign and target.id not in module_typevars ): data_name = target.id else: data_name = None if len(module_dict) > 0: ALIAS_DOCS_DICT[module_name] = module_dict if len(data_list) > 0: DATA_DICT[module_name] = data_list if module_typevars: TYPEVAR_DICT[module_name] = module_typevars return ALIAS_DOCS_DICT, DATA_DICT, TYPEVAR_DICT