Type _infer of binary operation nodes (#1659)

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

Author: DanielNoord

astroid/inference.py

@@ -510,18 +510,28 @@ def _infer_boolop(
 # UnaryOp, BinOp and AugAssign inferences
 
 
-def _filter_operation_errors(self, infer_callable, context, error):
+def _filter_operation_errors(
+    self: _T,
+    infer_callable: Callable[
+        [_T, InferenceContext | None],
+        Generator[InferenceResult | util.BadOperationMessage, None, None],
+    ],
+    context: InferenceContext | None,
+    error: type[util.BadOperationMessage],
+) -> Generator[InferenceResult, None, None]:
     for result in infer_callable(self, context):
         if isinstance(result, error):
             # For the sake of .infer(), we don't care about operation
             # errors, which is the job of pylint. So return something
             # which shows that we can't infer the result.
             yield util.Uninferable
         else:
-            yield result
+            yield result  # type: ignore[misc]
 
 
-def _infer_unaryop(self, context=None):
+def _infer_unaryop(
+    self: nodes.UnaryOp, context: InferenceContext | None = None
+) -> Generator[InferenceResult | util.BadUnaryOperationMessage, None, None]:
     """Infer what an UnaryOp should return when evaluated."""
     for operand in self.operand.infer(context):
         try:
@@ -579,16 +589,18 @@ def _infer_unaryop(self, context=None):
 
 @decorators.raise_if_nothing_inferred
 @decorators.path_wrapper
-def infer_unaryop(self, context=None):
+def infer_unaryop(
+    self: nodes.UnaryOp, context: InferenceContext | None = None, **kwargs: Any
+) -> Generator[InferenceResult, None, InferenceErrorInfo]:
     """Infer what an UnaryOp should return when evaluated."""
     yield from _filter_operation_errors(
         self, _infer_unaryop, context, util.BadUnaryOperationMessage
     )
-    return dict(node=self, context=context)
+    return InferenceErrorInfo(node=self, context=context)
 
 
 nodes.UnaryOp._infer_unaryop = _infer_unaryop
-nodes.UnaryOp._infer = infer_unaryop
+nodes.UnaryOp._infer = infer_unaryop  # type: ignore[assignment]
 
 
 def _is_not_implemented(const):
@@ -828,7 +840,9 @@ def _infer_binary_operation(left, right, binary_opnode, context, flow_factory):
     yield util.BadBinaryOperationMessage(left_type, binary_opnode.op, right_type)
 
 
-def _infer_binop(self, context):
+def _infer_binop(
+    self: nodes.BinOp, context: InferenceContext | None = None
+) -> Generator[InferenceResult | util.BadBinaryOperationMessage, None, None]:
     """Binary operation inference logic."""
     left = self.left
     right = self.right
@@ -855,14 +869,16 @@ def _infer_binop(self, context):
 
 @decorators.yes_if_nothing_inferred
 @decorators.path_wrapper
-def infer_binop(self, context=None):
+def infer_binop(
+    self: nodes.BinOp, context: InferenceContext | None = None, **kwargs: Any
+) -> Generator[InferenceResult, None, None]:
     return _filter_operation_errors(
         self, _infer_binop, context, util.BadBinaryOperationMessage
     )
 
 
 nodes.BinOp._infer_binop = _infer_binop
-nodes.BinOp._infer = infer_binop
+nodes.BinOp._infer = infer_binop  # type: ignore[assignment]
 
 COMPARE_OPS: dict[str, Callable[[Any, Any], bool]] = {
     "==": operator.eq,
@@ -932,8 +948,8 @@ def _do_compare(
 
 
 def _infer_compare(
-    self: nodes.Compare, context: InferenceContext | None = None
-) -> Iterator[nodes.Const | type[util.Uninferable]]:
+    self: nodes.Compare, context: InferenceContext | None = None, **kwargs: Any
+) -> Generator[nodes.Const | type[util.Uninferable], None, None]:
     """Chained comparison inference logic."""
     retval: bool | type[util.Uninferable] = True
 
@@ -961,7 +977,9 @@ def _infer_compare(
 nodes.Compare._infer = _infer_compare  # type: ignore[assignment]
 
 
-def _infer_augassign(self, context=None):
+def _infer_augassign(
+    self: nodes.AugAssign, context: InferenceContext | None = None
+) -> Generator[InferenceResult | util.BadBinaryOperationMessage, None, None]:
     """Inference logic for augmented binary operations."""
     context = context or InferenceContext()
 
@@ -989,14 +1007,16 @@ def _infer_augassign(self, context=None):
 
 @decorators.raise_if_nothing_inferred
 @decorators.path_wrapper
-def infer_augassign(self, context=None):
+def infer_augassign(
+    self: nodes.AugAssign, context: InferenceContext | None = None, **kwargs: Any
+) -> Generator[InferenceResult, None, None]:
     return _filter_operation_errors(
         self, _infer_augassign, context, util.BadBinaryOperationMessage
     )
 
 
 nodes.AugAssign._infer_augassign = _infer_augassign
-nodes.AugAssign._infer = infer_augassign
+nodes.AugAssign._infer = infer_augassign  # type: ignore[assignment]
 
 # End of binary operation inference.
 

astroid/nodes/node_classes.py

@@ -56,6 +56,7 @@ def _is_const(value):
 
 
 _NodesT = TypeVar("_NodesT", bound=NodeNG)
+_BadOpMessageT = TypeVar("_BadOpMessageT", bound=util.BadOperationMessage)
 
 AssignedStmtsPossibleNode = Union["List", "Tuple", "AssignName", "AssignAttr", None]
 AssignedStmtsCall = Callable[
@@ -71,6 +72,10 @@ def _is_const(value):
     [_NodesT, Optional[InferenceContext]],
     typing.Generator[InferenceResult, None, Optional[InferenceErrorInfo]],
 ]
+InferBinaryOperation = Callable[
+    [_NodesT, Optional[InferenceContext]],
+    typing.Generator[Union[InferenceResult, _BadOpMessageT], None, None],
+]
 
 
 @decorators.raise_if_nothing_inferred
@@ -1349,8 +1354,9 @@ def postinit(
     """
 
     # This is set by inference.py
-    def _infer_augassign(self, context=None):
-        raise NotImplementedError
+    _infer_augassign: ClassVar[
+        InferBinaryOperation[AugAssign, util.BadBinaryOperationMessage]
+    ]
 
     def type_errors(self, context=None):
         """Get a list of type errors which can occur during inference.
@@ -1449,8 +1455,7 @@ def postinit(self, left: NodeNG | None = None, right: NodeNG | None = None) -> N
         self.right = right
 
     # This is set by inference.py
-    def _infer_binop(self, context=None):
-        raise NotImplementedError
+    _infer_binop: ClassVar[InferBinaryOperation[BinOp, util.BadBinaryOperationMessage]]
 
     def type_errors(self, context=None):
         """Get a list of type errors which can occur during inference.
@@ -4232,8 +4237,9 @@ def postinit(self, operand: NodeNG | None = None) -> None:
         self.operand = operand
 
     # This is set by inference.py
-    def _infer_unaryop(self, context=None):
-        raise NotImplementedError
+    _infer_unaryop: ClassVar[
+        InferBinaryOperation[UnaryOp, util.BadUnaryOperationMessage]
+    ]
 
     def type_errors(self, context=None):
         """Get a list of type errors which can occur during inference.