3
3
from __future__ import annotations
4
4
5
5
import ast
6
+ import sys
6
7
from pathlib import Path
7
8
8
9
from typing_extensions import TypeAlias
@@ -158,27 +159,35 @@ def parse_module_attributes() -> tuple[AliasDocsDict, DataDict, TypeVarDict]:
158
159
inner_nodes = [node ]
159
160
160
161
for node in inner_nodes :
161
- # If we encounter an assignment annotated as "TypeAlias":
162
- if (
162
+ # Check if this node is a TypeAlias (type <name> = <value>)
163
+ # or an AnnAssign annotated as TypeAlias (<target>: TypeAlias = <value>).
164
+ is_type_alias = (
165
+ sys .version_info >= (3 , 12 ) and type (node ) is ast .TypeAlias
166
+ )
167
+ is_annotated_assignment_with_value = (
163
168
type (node ) is ast .AnnAssign
164
169
and type (node .annotation ) is ast .Name
165
170
and node .annotation .id == "TypeAlias"
166
171
and type (node .target ) is ast .Name
167
172
and node .value is not None
168
- ):
169
- alias_name = node .target .id
170
- def_node = node .value
171
- # If it's an Union, replace it with vertical bar notation
173
+ )
174
+ if is_type_alias or is_annotated_assignment_with_value :
175
+ alias_name = node .name .id if is_type_alias else node .target .id
176
+ definition_node = node .value
177
+
178
+ # If the definition is a Union, replace with vertical bar notation.
179
+ # Instead of "Union[Type1, Type2]", we'll have "Type1 | Type2".
172
180
if (
173
- type (def_node ) is ast .Subscript
174
- and type (def_node .value ) is ast .Name
175
- and def_node .value .id == "Union"
181
+ type (definition_node ) is ast .Subscript
182
+ and type (definition_node .value ) is ast .Name
183
+ and definition_node .value .id == "Union"
176
184
):
185
+ union_elements = definition_node .slice .elts
177
186
definition = " | " .join (
178
- ast .unparse (elem ) for elem in def_node . slice . elts
187
+ ast .unparse (elem ) for elem in union_elements
179
188
)
180
189
else :
181
- definition = ast .unparse (def_node )
190
+ definition = ast .unparse (definition_node )
182
191
183
192
definition = definition .replace ("npt." , "" )
184
193
if category_dict is None :
@@ -188,7 +197,7 @@ def parse_module_attributes() -> tuple[AliasDocsDict, DataDict, TypeVarDict]:
188
197
alias_info = category_dict [alias_name ]
189
198
continue
190
199
191
- # Check if it is a typing.TypeVar
200
+ # Check if it is a typing.TypeVar (<target> = TypeVar(...)).
192
201
elif (
193
202
type (node ) is ast .Assign
194
203
and type (node .targets [0 ]) is ast .Name
0 commit comments