handle binops in liveness checker fixes #157

core/ipyflow/analysis/live_refs.py

@@ -39,6 +39,17 @@
 _RESOLVER_EXCEPTIONS = ("get_ipython", "run_line_magic", "run_cell_magic")
 
 
+def _chain_root(node: ast.AST):
+    while True:
+        if isinstance(node, (ast.Attribute, ast.Subscript)):
+            node = node.value
+        elif isinstance(node, ast.Call):
+            node = node.func
+        else:
+            break
+    return node
+
+
 # TODO: have the logger warnings additionally raise exceptions for tests
 class ComputeLiveSymbolRefs(
     SaveOffAttributesMixin, SkipUnboundArgsMixin, VisitListsMixin, ast.NodeVisitor
@@ -329,7 +340,7 @@ def visit_Call(self, node: ast.Call) -> None:
             self.generic_visit(node.args)
             for kwarg in node.keywords:
                 self.visit(kwarg.value)
-        if not self._inside_attrsub:
+        if not self._inside_attrsub and not isinstance(_chain_root(node), ast.BinOp):
             self._add_attrsub_to_live_if_eligible(SymbolRef(node))
         with self.attrsub_context():
             self.visit(node.func)
@@ -359,13 +370,18 @@ def visit_Call(self, node: ast.Call) -> None:
             self.live |= call_live
 
     def visit_Attribute(self, node: ast.Attribute) -> None:
-        if not self._inside_attrsub:
+        if not self._inside_attrsub and not isinstance(_chain_root(node), ast.BinOp):
             self._add_attrsub_to_live_if_eligible(SymbolRef(node))
         with self.attrsub_context():
             self.visit(node.value)
 
+    def visit_BinOp(self, node: ast.BinOp) -> None:
+        with self.attrsub_context(False):
+            self.visit(node.left)
+            self.visit(node.right)
+
     def visit_Subscript(self, node: ast.Subscript) -> None:
-        if not self._inside_attrsub:
+        if not self._inside_attrsub and not isinstance(_chain_root(node), ast.BinOp):
             self._add_attrsub_to_live_if_eligible(SymbolRef(node))
         with self.attrsub_context():
             self.visit(node.value)

core/test/test_liveness_analysis.py

@@ -93,6 +93,26 @@ def test_subscript_is_live():
     assert live == {"foo", "bar", "baz"}
 
 
+def test_attribute_access_on_binop():
+    live, dead = compute_live_dead_symbol_refs("(a + b).c")
+    assert live == {"a", "b"}
+
+
+def test_call_on_binop():
+    live, dead = compute_live_dead_symbol_refs("(a + b).c()")
+    assert live == {"a", "b"}
+
+
+def test_attribute_reference_on_call_on_binop():
+    live, dead = compute_live_dead_symbol_refs("(a + b).c(d, e).f")
+    assert live == {"a", "b", "d", "e"}
+
+
+def test_attribute_reference_on_attribute_reference_on_binop():
+    live, dead = compute_live_dead_symbol_refs("(a + b).c.d")
+    assert live == {"a", "b"}
+
+
 def test_dict_literal():
     live, dead = compute_live_dead_symbol_refs("{'foo': bar}")
     assert live == {"bar"}

core/test/test_staleness_propagation.py

@@ -2763,3 +2763,14 @@ def baz():
         )
         run_cell("logging.info(y)")
         assert_detected()
+
+    @skipif_known_failing
+    def test_attribute_access_on_binop():
+        run_cell("a = 1")
+        run_cell("b = 2")
+        run_cell("c = (a + b).bit_length()")
+        run_cell("logging.info(c)")
+        assert_not_detected()
+        run_cell("a = 42")
+        run_cell("logging.info(c)")
+        assert_detected()