Add parameter typing to assigned_stmts (#1249)

Co-authored-by: Marc Mueller <[email protected]>

Author: DanielNoord

astroid/nodes/node_classes.py

@@ -42,7 +42,7 @@
 import typing
 import warnings
 from functools import lru_cache
-from typing import TYPE_CHECKING, Callable, Generator, Optional
+from typing import TYPE_CHECKING, Any, Callable, Generator, Optional, TypeVar, Union
 
 from astroid import decorators, mixins, util
 from astroid.bases import Instance, _infer_stmts
@@ -72,6 +72,20 @@ def _is_const(value):
     return isinstance(value, tuple(CONST_CLS))
 
 
+T_Nodes = TypeVar("T_Nodes", bound=NodeNG)
+
+AssignedStmtsPossibleNode = Union["List", "Tuple", "AssignName", "AssignAttr", None]
+AssignedStmtsCall = Callable[
+    [
+        T_Nodes,
+        AssignedStmtsPossibleNode,
+        Optional[InferenceContext],
+        Optional[typing.List[int]],
+    ],
+    Any,
+]
+
+
 @decorators.raise_if_nothing_inferred
 def unpack_infer(stmt, context=None):
     """recursively generate nodes inferred by the given statement.
@@ -672,6 +686,11 @@ def __init__(
             parent=parent,
         )
 
+    assigned_stmts: AssignedStmtsCall["AssignName"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
 
 class DelName(
     mixins.NoChildrenMixin, LookupMixIn, mixins.ParentAssignTypeMixin, NodeNG
@@ -993,6 +1012,11 @@ def postinit(
         if type_comment_posonlyargs is not None:
             self.type_comment_posonlyargs = type_comment_posonlyargs
 
+    assigned_stmts: AssignedStmtsCall["Arguments"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def _infer_name(self, frame, name):
         if self.parent is frame:
             return name
@@ -1246,6 +1270,11 @@ def postinit(self, expr: Optional[NodeNG] = None) -> None:
         """
         self.expr = expr
 
+    assigned_stmts: AssignedStmtsCall["AssignAttr"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def get_children(self):
         yield self.expr
 
@@ -1389,6 +1418,11 @@ def postinit(
         self.value = value
         self.type_annotation = type_annotation
 
+    assigned_stmts: AssignedStmtsCall["Assign"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def get_children(self):
         yield from self.targets
 
@@ -1481,6 +1515,11 @@ def postinit(
         self.value = value
         self.simple = simple
 
+    assigned_stmts: AssignedStmtsCall["AnnAssign"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def get_children(self):
         yield self.target
         yield self.annotation
@@ -1562,6 +1601,11 @@ def postinit(
         self.target = target
         self.value = value
 
+    assigned_stmts: AssignedStmtsCall["AugAssign"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     # This is set by inference.py
     def _infer_augassign(self, context=None):
         raise NotImplementedError
@@ -2028,6 +2072,11 @@ def postinit(
             self.ifs = ifs
         self.is_async = is_async
 
+    assigned_stmts: AssignedStmtsCall["Comprehension"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def assign_type(self):
         """The type of assignment that this node performs.
 
@@ -2715,6 +2764,11 @@ def __init__(
             parent=parent,
         )
 
+    assigned_stmts: AssignedStmtsCall["ExceptHandler"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def get_children(self):
         if self.type is not None:
             yield self.type
@@ -2873,6 +2927,11 @@ def postinit(
             self.orelse = orelse
         self.type_annotation = type_annotation
 
+    assigned_stmts: AssignedStmtsCall["For"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     @decorators.cachedproperty
     def blockstart_tolineno(self):
         """The line on which the beginning of this block ends.
@@ -3573,6 +3632,11 @@ def __init__(
             parent=parent,
         )
 
+    assigned_stmts: AssignedStmtsCall["List"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def pytype(self):
         """Get the name of the type that this node represents.
 
@@ -3999,6 +4063,11 @@ def postinit(self, value: Optional[NodeNG] = None) -> None:
         """
         self.value = value
 
+    assigned_stmts: AssignedStmtsCall["Starred"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def get_children(self):
         yield self.value
 
@@ -4325,6 +4394,11 @@ def __init__(
             parent=parent,
         )
 
+    assigned_stmts: AssignedStmtsCall["Tuple"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def pytype(self):
         """Get the name of the type that this node represents.
 
@@ -4619,6 +4693,11 @@ def postinit(
             self.body = body
         self.type_annotation = type_annotation
 
+    assigned_stmts: AssignedStmtsCall["With"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     @decorators.cachedproperty
     def blockstart_tolineno(self):
         """The line on which the beginning of this block ends.
@@ -4922,6 +5001,11 @@ def postinit(self, target: NodeNG, value: NodeNG) -> None:
         self.target = target
         self.value = value
 
+    assigned_stmts: AssignedStmtsCall["NamedExpr"]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
     def frame(self):
         """The first parent frame node.
 
@@ -5291,6 +5375,9 @@ def postinit(
         ],
         Generator[NodeNG, None, None],
     ]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
 
 
 class MatchClass(Pattern):
@@ -5393,6 +5480,9 @@ def postinit(self, *, name: Optional[AssignName]) -> None:
         ],
         Generator[NodeNG, None, None],
     ]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
 
 
 class MatchAs(mixins.AssignTypeMixin, Pattern):
@@ -5459,6 +5549,9 @@ def postinit(
         ],
         Generator[NodeNG, None, None],
     ]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
 
 
 class MatchOr(Pattern):

astroid/protocols.py

@@ -32,7 +32,7 @@
 import itertools
 import operator as operator_mod
 import sys
-from typing import Generator, Optional
+from typing import Any, Generator, List, Optional, Union
 
 from astroid import arguments, bases, decorators, helpers, nodes, util
 from astroid.const import Context
@@ -274,7 +274,12 @@ def _resolve_looppart(parts, assign_path, context):
 
 
 @decorators.raise_if_nothing_inferred
-def for_assigned_stmts(self, node=None, context=None, assign_path=None):
+def for_assigned_stmts(
+    self: Union[nodes.For, nodes.Comprehension],
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     if isinstance(self, nodes.AsyncFor) or getattr(self, "is_async", False):
         # Skip inferring of async code for now
         return dict(node=self, unknown=node, assign_path=assign_path, context=context)
@@ -291,7 +296,12 @@ def for_assigned_stmts(self, node=None, context=None, assign_path=None):
 nodes.Comprehension.assigned_stmts = for_assigned_stmts
 
 
-def sequence_assigned_stmts(self, node=None, context=None, assign_path=None):
+def sequence_assigned_stmts(
+    self: Union[nodes.Tuple, nodes.List],
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     if assign_path is None:
         assign_path = []
     try:
@@ -314,7 +324,12 @@ def sequence_assigned_stmts(self, node=None, context=None, assign_path=None):
 nodes.List.assigned_stmts = sequence_assigned_stmts
 
 
-def assend_assigned_stmts(self, node=None, context=None, assign_path=None):
+def assend_assigned_stmts(
+    self: Union[nodes.AssignName, nodes.AssignAttr],
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     return self.parent.assigned_stmts(node=self, context=context)
 
 
@@ -381,7 +396,12 @@ def _arguments_infer_argname(self, name, context):
         yield util.Uninferable
 
 
-def arguments_assigned_stmts(self, node=None, context=None, assign_path=None):
+def arguments_assigned_stmts(
+    self: nodes.Arguments,
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     if context.callcontext:
         callee = context.callcontext.callee
         while hasattr(callee, "_proxied"):
@@ -406,7 +426,12 @@ def arguments_assigned_stmts(self, node=None, context=None, assign_path=None):
 
 
 @decorators.raise_if_nothing_inferred
-def assign_assigned_stmts(self, node=None, context=None, assign_path=None):
+def assign_assigned_stmts(
+    self: Union[nodes.AugAssign, nodes.Assign, nodes.AnnAssign],
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     if not assign_path:
         yield self.value
         return None
@@ -417,7 +442,12 @@ def assign_assigned_stmts(self, node=None, context=None, assign_path=None):
     return dict(node=self, unknown=node, assign_path=assign_path, context=context)
 
 
-def assign_annassigned_stmts(self, node=None, context=None, assign_path=None):
+def assign_annassigned_stmts(
+    self: nodes.AnnAssign,
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     for inferred in assign_assigned_stmts(self, node, context, assign_path):
         if inferred is None:
             yield util.Uninferable
@@ -471,7 +501,12 @@ def _resolve_assignment_parts(parts, assign_path, context):
 
 
 @decorators.raise_if_nothing_inferred
-def excepthandler_assigned_stmts(self, node=None, context=None, assign_path=None):
+def excepthandler_assigned_stmts(
+    self: nodes.ExceptHandler,
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     for assigned in node_classes.unpack_infer(self.type):
         if isinstance(assigned, nodes.ClassDef):
             assigned = objects.ExceptionInstance(assigned)
@@ -522,7 +557,12 @@ def _infer_context_manager(self, mgr, context):
 
 
 @decorators.raise_if_nothing_inferred
-def with_assigned_stmts(self, node=None, context=None, assign_path=None):
+def with_assigned_stmts(
+    self: nodes.With,
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     """Infer names and other nodes from a *with* statement.
 
     This enables only inference for name binding in a *with* statement.
@@ -595,7 +635,12 @@ def __enter__(self):
 
 
 @decorators.raise_if_nothing_inferred
-def named_expr_assigned_stmts(self, node, context=None, assign_path=None):
+def named_expr_assigned_stmts(
+    self: nodes.NamedExpr,
+    node: node_classes.AssignedStmtsPossibleNode,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     """Infer names and other nodes from an assignment expression"""
     if self.target == node:
         yield from self.value.infer(context=context)
@@ -612,7 +657,12 @@ def named_expr_assigned_stmts(self, node, context=None, assign_path=None):
 
 
 @decorators.yes_if_nothing_inferred
-def starred_assigned_stmts(self, node=None, context=None, assign_path=None):
+def starred_assigned_stmts(
+    self: nodes.Starred,
+    node: node_classes.AssignedStmtsPossibleNode = None,
+    context: Optional[InferenceContext] = None,
+    assign_path: Optional[List[int]] = None,
+) -> Any:
     """
     Arguments:
         self: nodes.Starred