Implement equality constraint (#2927)

Author: zenlyj

ChangeLog

@@ -7,6 +7,10 @@ What's New in astroid 4.1.0?
 ============================
 Release date: TBA
 
+* Add support for equality constraints (``==``, ``!=``) in inference.
+  Closes pylint-dev/pylint#3632
+  Closes pylint-dev/pylint#3633
+
 * Ensure ``ast.JoinedStr`` nodes are ``Uninferable`` when the ``ast.FormattedValue`` is
   ``Uninferable``. This prevents ``unexpected-keyword-arg`` messages in Pylint
   where the ``Uninferable`` string appeared in function arguments that were

astroid/constraint.py

@@ -174,6 +174,62 @@ def satisfied_by(self, inferred: InferenceResult) -> bool:
             return True
 
 
+class EqualityConstraint(Constraint):
+    """Represents a "==" or "!=" constraint."""
+
+    def __init__(self, node: nodes.NodeNG, operand: nodes.NodeNG, negate: bool) -> None:
+        super().__init__(node=node, negate=negate)
+        self.operand = operand
+
+    @classmethod
+    def match(
+        cls, node: _NameNodes, expr: nodes.NodeNG, negate: bool = False
+    ) -> Self | None:
+        """Return a new constraint for node if expr matches one of these patterns:
+
+        - "node == operand" or "operand == node": use given negate value
+        - "node != operand" or "operand != node": flip negate value
+
+        Return None if no pattern matches.
+        """
+        if isinstance(expr, nodes.Compare) and len(expr.ops) == 1:
+            left = expr.left
+            op, right = expr.ops[0]
+            matches_left = _matches(left, node)
+
+            if op in {"==", "!="} and (matches_left or _matches(right, node)):
+                operand = right if matches_left else left
+                negate = (op == "==" and negate) or (op == "!=" and not negate)
+                return cls(node=node, operand=operand, negate=negate)
+
+        return None
+
+    def satisfied_by(self, inferred: InferenceResult) -> bool:
+        """Return True for uninferable/ambiguous results, or depending on negate flag:
+
+        - negate=False: satisfied when both operands are equal.
+        - negate=True: satisfied when both operands are not equal.
+
+        Only comparisons between constants and callables are supported.
+        """
+        if inferred is util.Uninferable:
+            return True
+
+        operand_inferred = util.safe_infer(self.operand)
+        if operand_inferred is util.Uninferable or operand_inferred is None:
+            return True
+
+        if isinstance(inferred, nodes.Const) and isinstance(
+            operand_inferred, nodes.Const
+        ):
+            return self.negate ^ (inferred.value == operand_inferred.value)
+
+        if inferred.callable() and operand_inferred.callable():
+            return self.negate ^ (inferred is operand_inferred)
+
+        return True
+
+
 def get_constraints(
     expr: _NameNodes, frame: nodes.LocalsDictNodeNG
 ) -> dict[nodes.If | nodes.IfExp, set[Constraint]]:
@@ -209,6 +265,7 @@ def get_constraints(
         NoneConstraint,
         BooleanConstraint,
         TypeConstraint,
+        EqualityConstraint,
     )
 )
 """All supported constraint types."""

tests/test_constraint.py

@@ -15,6 +15,10 @@
 from astroid.util import Uninferable
 
 
+def node_info(node: nodes.NodeNG) -> str:
+    return f"Inference of {node.as_string()!r} at line {node.lineno}"
+
+
 def common_params(node: str) -> pytest.MarkDecorator:
     return pytest.mark.parametrize(
         ("condition", "satisfy_val", "fail_val"),
@@ -25,6 +29,10 @@ def common_params(node: str) -> pytest.MarkDecorator:
             (f"not {node}", None, 3),
             (f"isinstance({node}, int)", 3, None),
             (f"isinstance({node}, (int, str))", 3, None),
+            (f"{node} == 3", 3, None),
+            (f"{node} != 3", None, 3),
+            (f"3 == {node}", 3, None),
+            (f"3 != {node}", None, 3),
         ),
     )
 
@@ -267,17 +275,19 @@ def f2(y, x = {satisfy_val}):
             )
     """)
     for node in (node1, node2):
+        msg = node_info(node)
         inferred = node.inferred()
-        assert len(inferred) == 2
-        assert isinstance(inferred[0], nodes.Const)
-        assert inferred[0].value == fail_val
+        assert len(inferred) == 2, msg
+        assert isinstance(inferred[0], nodes.Const), msg
+        assert inferred[0].value == fail_val, msg
 
-        assert inferred[1] is Uninferable
+        assert inferred[1] is Uninferable, msg
 
     for node in (node3, node4):
+        msg = node_info(node)
         inferred = node.inferred()
-        assert len(inferred) == 1
-        assert inferred[0] is Uninferable
+        assert len(inferred) == 1, msg
+        assert inferred[0] is Uninferable, msg
 
 
 @common_params(node="x")
@@ -839,11 +849,12 @@ class C(A, B):
     """)
 
     for node in (n1, n2, n3):
+        msg = node_info(node)
         inferred = node.inferred()
-        assert len(inferred) == 1
-        assert isinstance(inferred[0], Instance)
-        assert isinstance(inferred[0]._proxied, nodes.ClassDef)
-        assert inferred[0].name == "C"
+        assert len(inferred) == 1, msg
+        assert isinstance(inferred[0], Instance), msg
+        assert isinstance(inferred[0]._proxied, nodes.ClassDef), msg
+        assert inferred[0].name == "C", msg
 
 
 def test_isinstance_diamond_inheritance():
@@ -879,11 +890,12 @@ class D(B, C):
     """)
 
     for node in (n1, n2, n3, n4):
+        msg = node_info(node)
         inferred = node.inferred()
-        assert len(inferred) == 1
-        assert isinstance(inferred[0], Instance)
-        assert isinstance(inferred[0]._proxied, nodes.ClassDef)
-        assert inferred[0].name == "D"
+        assert len(inferred) == 1, msg
+        assert isinstance(inferred[0], Instance), msg
+        assert isinstance(inferred[0]._proxied, nodes.ClassDef), msg
+        assert inferred[0].name == "D", msg
 
 
 def test_isinstance_keyword_arguments():
@@ -901,10 +913,11 @@ def test_isinstance_keyword_arguments():
     """)
 
     for node in (n1, n2):
+        msg = node_info(node)
         inferred = node.inferred()
-        assert len(inferred) == 1
-        assert isinstance(inferred[0], nodes.Const)
-        assert inferred[0].value == 3
+        assert len(inferred) == 1, msg
+        assert isinstance(inferred[0], nodes.Const), msg
+        assert inferred[0].value == 3, msg
 
 
 def test_isinstance_extra_argument():
@@ -999,3 +1012,132 @@ def test_isinstance_uninferable():
         assert len(inferred) == 1
         assert isinstance(inferred[0], nodes.Const)
         assert inferred[0].value == 3
+
+
+def test_equality_callable():
+    """Test constraint for equality of callables."""
+    node1, node2, node3, node4, node5, node6 = builder.extract_node("""
+    class Foo:
+        pass
+
+    def bar():
+        pass
+
+    baz = lambda i : i
+
+    x, y, z = Foo, bar, baz
+
+    if x == Foo:
+        x  #@
+    if x != Foo:
+        x  #@
+
+    if y == bar:
+        y  #@
+    if y != bar:
+        y  #@
+
+    if z == baz:
+        z  #@
+    if z != baz:
+        z  #@
+    """)
+
+    inferred = node1.inferred()
+    assert len(inferred) == 1
+    assert isinstance(inferred[0], nodes.ClassDef)
+    assert inferred[0].name == "Foo"
+
+    inferred = node3.inferred()
+    assert len(inferred) == 1
+    assert isinstance(inferred[0], nodes.FunctionDef)
+    assert inferred[0].name == "bar"
+
+    inferred = node5.inferred()
+    assert len(inferred) == 1
+    assert isinstance(inferred[0], nodes.Lambda)
+
+    for node in (node2, node4, node6):
+        msg = node_info(node)
+        inferred = node.inferred()
+        assert len(inferred) == 1, msg
+        assert inferred[0] is Uninferable, msg
+
+
+def test_equality_uninferable_operand():
+    """Test that equality constraint is satisfied when either operand is uninferable."""
+    node1, node2, node3, node4 = builder.extract_node("""
+    def f1(x):
+        if x == 3:
+            x  #@
+
+        if x != 3:
+            x  #@
+
+    def f2(y):
+        x = 3
+        if x == y:
+            x  #@
+
+        if x != y:
+            x  #@
+    """)
+
+    for node in (node1, node2):
+        msg = node_info(node)
+        inferred = node.inferred()
+        assert len(inferred) == 1, msg
+        assert inferred[0] is Uninferable, msg
+
+    for node in (node3, node4):
+        msg = node_info(node)
+        inferred = node.inferred()
+        assert len(inferred) == 1, msg
+        assert isinstance(inferred[0], nodes.Const), msg
+        assert inferred[0].value == 3, msg
+
+
+def test_equality_ambiguous_operand():
+    """Test that equality constraint is satisfied when the compared operand has multiple inferred values."""
+    node1, node2 = builder.extract_node("""
+    def f(y = 1):
+        x = 3
+        if x == y:
+            x  #@
+
+        if x != y:
+            x  #@
+    """)
+
+    for node in (node1, node2):
+        msg = node_info(node)
+        inferred = node.inferred()
+        assert len(inferred) == 1, msg
+        assert isinstance(inferred[0], nodes.Const), msg
+        assert inferred[0].value == 3, msg
+
+
+def test_equality_fractions():
+    """Test that equality constraint is satisfied when both operands are fractions."""
+    node1, node2, node3, node4 = builder.extract_node("""
+    from fractions import Fraction
+
+    x = Fraction(1, 3)
+    y = Fraction(1, 3)
+
+    if x == y:
+        x  #@
+        y  #@
+
+    if x != y:
+        x  #@
+        y  #@
+    """)
+
+    for node in (node1, node2, node3, node4):
+        msg = node_info(node)
+        inferred = node.inferred()
+        assert len(inferred) == 1, msg
+        assert isinstance(inferred[0], Instance), msg
+        assert isinstance(inferred[0]._proxied, nodes.ClassDef), msg
+        assert inferred[0]._proxied.name == "Fraction", msg