Skip to content

Commit 047ac68

Browse files
JasonGrace2282JasonGrace2282
committed
fix(autoaliasattr): search for type aliases under if TYPE_CHECKING
1 parent f8e3408 commit 047ac68

File tree

1 file changed

+77
-42
lines changed

1 file changed

+77
-42
lines changed

‎manim/utils/docbuild/module_parsing.py

Lines changed: 77 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@
4444

4545
MANIM_ROOT = Path(__file__).resolve().parent.parent.parent
4646

47+
# In the following, we will use ``type(xyz) is xyz_type`` instead of
48+
# isinstance checks to make sure no subclasses of the type pass the
49+
# check
50+
4751

4852
def parse_module_attributes() -> tuple[AliasDocsDict, DataDict]:
4953
"""Read all files, generate Abstract Syntax Trees from them, and
@@ -70,7 +74,7 @@ def parse_module_attributes() -> tuple[AliasDocsDict, DataDict]:
7074
for module_path in MANIM_ROOT.rglob("*.py"):
7175
module_name = module_path.resolve().relative_to(MANIM_ROOT)
7276
module_name = list(module_name.parts)
73-
module_name[-1] = module_name[-1][:-3] # remove .py
77+
module_name[-1] = module_name[-1].removesuffix(".py")
7478
module_name = ".".join(module_name)
7579

7680
module_content = module_path.read_text(encoding="utf-8")
@@ -107,53 +111,84 @@ def parse_module_attributes() -> tuple[AliasDocsDict, DataDict]:
107111
data_list.append(data_name)
108112
continue
109113

110-
# If we encounter an assignment annotated as "TypeAlias":
114+
# if it's defined under if TYPE_CHECKING
115+
# go through the body of the if statement
111116
if (
112-
type(node) is ast.AnnAssign
113-
and type(node.annotation) is ast.Name
114-
and node.annotation.id == "TypeAlias"
115-
and type(node.target) is ast.Name
116-
and node.value is not None
117+
# NOTE: This logic does not (and cannot)
118+
# check if the comparison is against a
119+
# variable called TYPE_CHECKING
120+
# It also says that you cannot do the following
121+
# import typing as foo
122+
# if foo.TYPE_CHECKING:
123+
# BAR: TypeAlias = ...
124+
type(node) is ast.If
125+
and (
126+
(
127+
# if TYPE_CHECKING
128+
type(node.test) is ast.Name
129+
and node.test.id == "TYPE_CHECKING"
130+
)
131+
or (
132+
# if typing.TYPE_CHECKING
133+
type(node.test) is ast.Attribute
134+
and type(node.test.value) is ast.Name
135+
and node.test.value.id == "typing"
136+
and node.test.attr == "TYPE_CHECKING"
137+
)
138+
)
117139
):
118-
alias_name = node.target.id
119-
def_node = node.value
120-
# If it's an Union, replace it with vertical bar notation
140+
inner_nodes = node.body
141+
else:
142+
inner_nodes = [node]
143+
144+
for node in inner_nodes:
145+
# If we encounter an assignment annotated as "TypeAlias":
121146
if (
122-
type(def_node) is ast.Subscript
123-
and type(def_node.value) is ast.Name
124-
and def_node.value.id == "Union"
147+
type(node) is ast.AnnAssign
148+
and type(node.annotation) is ast.Name
149+
and node.annotation.id == "TypeAlias"
150+
and type(node.target) is ast.Name
151+
and node.value is not None
125152
):
126-
definition = " | ".join(
127-
ast.unparse(elem) for elem in def_node.slice.elts
128-
)
153+
alias_name = node.target.id
154+
def_node = node.value
155+
# If it's an Union, replace it with vertical bar notation
156+
if (
157+
type(def_node) is ast.Subscript
158+
and type(def_node.value) is ast.Name
159+
and def_node.value.id == "Union"
160+
):
161+
definition = " | ".join(
162+
ast.unparse(elem) for elem in def_node.slice.elts
163+
)
164+
else:
165+
definition = ast.unparse(def_node)
166+
167+
definition = definition.replace("npt.", "")
168+
if category_dict is None:
169+
module_dict[""] = {}
170+
category_dict = module_dict[""]
171+
category_dict[alias_name] = {"definition": definition}
172+
alias_info = category_dict[alias_name]
173+
continue
174+
175+
# If here, the node is not a TypeAlias definition
176+
alias_info = None
177+
178+
# It could still be a module attribute definition.
179+
# Does the assignment have a target of type Name? Then
180+
# it could be considered a definition of a module attribute.
181+
if type(node) is ast.AnnAssign:
182+
target = node.target
183+
elif type(node) is ast.Assign and len(node.targets) == 1:
184+
target = node.targets[0]
129185
else:
130-
definition = ast.unparse(def_node)
131-
132-
definition = definition.replace("npt.", "")
133-
if category_dict is None:
134-
module_dict[""] = {}
135-
category_dict = module_dict[""]
136-
category_dict[alias_name] = {"definition": definition}
137-
alias_info = category_dict[alias_name]
138-
continue
139-
140-
# If here, the node is not a TypeAlias definition
141-
alias_info = None
142-
143-
# It could still be a module attribute definition.
144-
# Does the assignment have a target of type Name? Then
145-
# it could be considered a definition of a module attribute.
146-
if type(node) is ast.AnnAssign:
147-
target = node.target
148-
elif type(node) is ast.Assign and len(node.targets) == 1:
149-
target = node.targets[0]
150-
else:
151-
target = None
186+
target = None
152187

153-
if type(target) is ast.Name:
154-
data_name = target.id
155-
else:
156-
data_name = None
188+
if type(target) is ast.Name:
189+
data_name = target.id
190+
else:
191+
data_name = None
157192

158193
if len(module_dict) > 0:
159194
ALIAS_DOCS_DICT[module_name] = module_dict

0 commit comments

Comments
 (0)