Implement `And` and `Or` constraints (#2977)
Co-authored-by: Pierre Sassoulas <pierre.sassoulas@gmail.com>
diff --git a/astroid/constraint.py b/astroid/constraint.py
index 8f7e6b2..af6c59f 100644
--- a/astroid/constraint.py
+++ b/astroid/constraint.py
@@ -247,6 +247,61 @@
return True
+class _CompoundConstraint(Constraint):
+ """Represents an "x and y" or "x or y" constraint."""
+
+ def __init__(
+ self,
+ node: nodes.NodeNG,
+ op: str,
+ children: list[Constraint],
+ negate: bool,
+ ) -> None:
+ super().__init__(node=node, negate=negate)
+ self.op = op
+ self.children = children
+
+ @classmethod
+ def match(
+ cls, node: _NameNodes, expr: nodes.NodeNG, negate: bool = False
+ ) -> Self | None:
+ """Return a new constraint for node if expr matches an "x and y" or
+ "x or y" pattern.
+
+ Return None if expr is not a supported boolean expression, or if any
+ operand does not match a constraint pattern.
+ """
+ if not (isinstance(expr, nodes.BoolOp) and expr.op in {"and", "or"}):
+ return None
+
+ children: list[Constraint] = []
+ for value in expr.values:
+ matches = list(_match_constraint(node, value, negate))
+ if not matches:
+ return None
+ children.extend(matches)
+
+ return cls(node=node, op=expr.op, children=children, negate=negate)
+
+ def satisfied_by(
+ self, inferred: InferenceResult, context: InferenceContext
+ ) -> bool:
+ """Return True for uninferable results, or depending on op and negate:
+
+ - negate=False: all children must be satisfied for "and", or any for "or".
+ - negate=True: any child must be satisfied for "and", or all for "or".
+ """
+ if inferred is util.Uninferable:
+ return True
+
+ results = (
+ constraint.satisfied_by(inferred, context) for constraint in self.children
+ )
+
+ strict = (self.op == "and") ^ self.negate
+ return all(results) if strict else any(results)
+
+
def get_constraints(
expr: _NameNodes, frame: nodes.LocalsDictNodeNG
) -> dict[nodes.NodeNG, set[Constraint]]:
@@ -310,15 +365,15 @@
constraints_mapping[if_expr] = constraints
-ALL_CONSTRAINT_CLASSES = frozenset(
- (
- NoneConstraint,
- BooleanConstraint,
- TypeConstraint,
- EqualityConstraint,
- )
-)
-"""All supported constraint types."""
+_CONSTRAINTS_BY_NODE_TYPE: dict[type[nodes.NodeNG], tuple[type[Constraint], ...]] = {
+ nodes.Attribute: (BooleanConstraint,),
+ nodes.BoolOp: (_CompoundConstraint,),
+ nodes.Call: (TypeConstraint,),
+ nodes.Compare: (NoneConstraint, EqualityConstraint),
+ nodes.Name: (BooleanConstraint,),
+ nodes.UnaryOp: (BooleanConstraint,),
+}
+"""Constraint types that can match each expression node type."""
def _matches(node1: nodes.NodeNG | bases.Proxy, node2: nodes.NodeNG) -> bool:
@@ -337,7 +392,7 @@
node: _NameNodes, expr: nodes.NodeNG, invert: bool = False
) -> Iterator[Constraint]:
"""Yields all constraint patterns for node that match."""
- for constraint_cls in ALL_CONSTRAINT_CLASSES:
+ for constraint_cls in _CONSTRAINTS_BY_NODE_TYPE.get(type(expr), ()):
constraint = constraint_cls.match(node, expr, invert)
if constraint:
yield constraint
diff --git a/doc/whatsnew/fragments/2977.feature b/doc/whatsnew/fragments/2977.feature
new file mode 100644
index 0000000..c21f46c
--- /dev/null
+++ b/doc/whatsnew/fragments/2977.feature
@@ -0,0 +1,8 @@
+Inference now narrows values through conditions joined by ``and`` and ``or``.
+A variable tested by several constraints at once, as in
+``if isinstance(apple, int) and apple != 3``, keeps only the values satisfying
+every operand, and the ``or`` and negated forms are handled too. Inference stays
+conservative when an operand cannot be inferred or does not constrain the
+variable.
+
+Closes #2977
diff --git a/tests/test_constraint.py b/tests/test_constraint.py
index b77b3b5..22cd1d7 100644
--- a/tests/test_constraint.py
+++ b/tests/test_constraint.py
@@ -33,6 +33,8 @@
(f"{node} != 3", None, 3),
(f"3 == {node}", 3, None),
(f"3 != {node}", None, 3),
+ (f"isinstance({node}, int) and {node} == 3", 3, 5),
+ (f"isinstance({node}, str) or {node} == 3", 3, None),
),
)
@@ -1316,3 +1318,148 @@
assert len(inferred) == 1
assert isinstance(inferred[0], nodes.Const)
assert inferred[0].value == 1
+
+
+@pytest.mark.parametrize(
+ ("condition", "satisfy_val", "fail_val"),
+ (
+ pytest.param(
+ "x is not None and (isinstance(x, int) and x == 3)",
+ 3,
+ 5,
+ id="nested-and",
+ ),
+ pytest.param(
+ "x is not None and x and isinstance(x, int) and x == 3",
+ 3,
+ 0,
+ id="and-with-multiple-operands",
+ ),
+ pytest.param(
+ "x is None or (isinstance(x, str) or x == 3)",
+ 3,
+ 5,
+ id="nested-or",
+ ),
+ pytest.param(
+ "x is None or not x or isinstance(x, str) or x == 3",
+ 0,
+ 5,
+ id="or-with-multiple-operands",
+ ),
+ pytest.param(
+ "x is not None and (isinstance(x, bool) or x == 3)",
+ True,
+ 5,
+ id="and-with-nested-or",
+ ),
+ pytest.param(
+ "x is None or (isinstance(x, int) and x == 3)",
+ 3,
+ 5,
+ id="or-with-nested-and",
+ ),
+ pytest.param(
+ "x == 3 or isinstance(x, int) and x == 5",
+ 3,
+ None,
+ id="and-precedence-over-or",
+ ),
+ ),
+)
+def test_compound_constraint(
+ condition: str, satisfy_val: int | None, fail_val: int | None
+) -> None:
+ """Test compound constraints in both the body and else branch of an if."""
+ n1, n2, n3, n4 = builder.extract_node(f"""
+ def f1(x = {fail_val}):
+ if {condition}:
+ x #@
+ else:
+ x #@
+
+ def f2(x = {satisfy_val}):
+ if {condition}:
+ x #@
+ else:
+ x #@
+ """)
+
+ for node in (n1, n4):
+ msg = node_info(node)
+ inferred = node.inferred()
+ assert len(inferred) == 1, msg
+ assert inferred[0] is Uninferable, msg
+
+ for node, expected in ((n2, fail_val), (n3, satisfy_val)):
+ msg = node_info(node)
+ inferred = node.inferred()
+ assert len(inferred) == 2, msg
+ assert isinstance(inferred[0], nodes.Const), msg
+ assert inferred[0].value == expected, msg
+ assert inferred[1] is Uninferable, msg
+
+
+@pytest.mark.parametrize(
+ "condition",
+ (
+ pytest.param("isinstance(x, classinfo) and x == 3", id="and"),
+ pytest.param("isinstance(x, classinfo) or x == 5", id="or"),
+ ),
+)
+def test_compound_constraint_with_uninferable_child(condition: str) -> None:
+ """Test that inference remains conservative when a child constraint is unknown."""
+ n1, n2 = builder.extract_node(f"""
+ def f(classinfo, x = 3):
+ if {condition}:
+ x #@
+ else:
+ x #@
+ """)
+
+ for node in (n1, n2):
+ msg = node_info(node)
+ inferred = node.inferred()
+ assert len(inferred) == 2, msg
+ assert isinstance(inferred[0], nodes.Const), msg
+ assert inferred[0].value == 3, msg
+ assert inferred[1].value is Uninferable, msg
+
+
+@pytest.mark.parametrize(
+ ("condition", "value"),
+ (
+ pytest.param("not x and y", 3, id="and"),
+ pytest.param(
+ "x is not None and (not x and y)",
+ 3,
+ id="nested-and-inner",
+ ),
+ pytest.param(
+ "x is not None and (not x and y)",
+ None,
+ id="nested-and-outer",
+ ),
+ pytest.param("not x or y", 3, id="or"),
+ pytest.param(
+ "x is None or (not x or y)",
+ 3,
+ id="nested-or",
+ ),
+ ),
+)
+def test_expression_with_non_constraint_is_ignored(
+ condition: str, value: int | None
+) -> None:
+ """Test that an expression is ignored when an operand does not constrain the target variable."""
+ node = builder.extract_node(f"""
+ x, y = {value}, None
+
+ if {condition}:
+ x #@
+ """)
+
+ inferred = node.inferred()
+ assert len(inferred) == 1
+ assert isinstance(inferred[0], nodes.Const)
+ assert inferred[0].value == value
