Source code for manim.utils.docbuild.module_parsing
"""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
[docs]
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 <name> = <value>)
# or an AnnAssign annotated as TypeAlias (<target>: TypeAlias = <value>).
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 (<target> = 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