Add TypeAlias.assigned_stmts() (#2252)

pylint-dev · Jul 16, 2023 · 623482b · 623482b
1 parent 4c3bebb
commit 623482b
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/astroid/nodes/node_classes.py b/astroid/nodes/node_classes.py
@@ -19,6 +19,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    ClassVar,
     Literal,
     Optional,
     Union,
@@ -4058,6 +4059,18 @@ def _infer(
     ) -> Iterator[TypeAlias]:
         yield self
 
+    assigned_stmts: ClassVar[
+        Callable[
+            [
+                TypeAlias,
+                AssignName,
+                InferenceContext | None,
+                None,
+            ],
+            Generator[NodeNG, None, None],
+        ]
+    ] = protocols.assign_assigned_stmts
+
 
 class TypeVar(_base_nodes.AssignTypeNode):
     """Class representing a :class:`ast.TypeVar` node.

diff --git a/astroid/protocols.py b/astroid/protocols.py
@@ -441,7 +441,7 @@ def arguments_assigned_stmts(
 
 @decorators.raise_if_nothing_inferred
 def assign_assigned_stmts(
-    self: nodes.AugAssign | nodes.Assign | nodes.AnnAssign,
+    self: nodes.AugAssign | nodes.Assign | nodes.AnnAssign | nodes.TypeAlias,
     node: node_classes.AssignedStmtsPossibleNode = None,
     context: InferenceContext | None = None,
     assign_path: list[int] | None = None,

diff --git a/tests/test_type_params.py b/tests/test_type_params.py
@@ -37,6 +37,9 @@ def test_type_alias() -> None:
 
     assert node.statement() is node
 
+    assigned = next(node.assigned_stmts())
+    assert assigned is node.value
+
 
 def test_type_param_spec() -> None:
     node = extract_node("type Alias[**P] = Callable[P, int]")