Merge pull request #51 from jakkdl/split_up_assert

src/shed/_codemods.py

@@ -401,3 +401,18 @@ def collapse_isinstance_checks(self, _, updated_node):
             )
             return updated_node.left.with_changes(args=[left_target, merged_type])
         return updated_node
+
+    # helper function to handle recursive structures created by e.g. `assert a and b and c`
+    def _flatten_bool(self, expr):
+        if m.matches(expr, m.BooleanOperation(operator=m.And())):
+            return self._flatten_bool(expr.left) + self._flatten_bool(expr.right)
+        return [expr]
+
+    # split `assert a and b` into `assert a` and `assert b`
+    @m.leave(m.Assert(test=m.BooleanOperation(operator=m.And())))
+    def split_assert_and(self, _, updated_node):
+        flat_nodes = self._flatten_bool(updated_node.test)
+        nodes = [updated_node.with_changes(test=flat_nodes[0])]
+        for node in flat_nodes[1:]:
+            nodes.append(cst.Assert(node))
+        return cst.FlattenSentinel(nodes)

tests/recorded/split_assert.txt

@@ -0,0 +1,62 @@
+assert a and b # one
+
+assert a and b and c # two
+
+assert (a and b) # three
+
+assert a(1, 2) and b(3, 7) # four
+
+def f():
+    assert a and b
+
+assert (a and b) and (c and d)
+
+assert a or b and c
+assert a and b or c
+
+assert a and (b or c)
+
+assert (a or b) and c
+
+assert a or b
+
+assert a
+
+================================================================================
+
+assert a
+assert b  # one
+
+assert a
+assert b
+assert c  # two
+
+assert a
+assert b  # three
+
+assert a(1, 2)
+assert b(3, 7)  # four
+
+
+def f():
+    assert a
+    assert b
+
+
+assert a
+assert b
+assert c
+assert d
+
+assert a or b and c
+assert a and b or c
+
+assert a
+assert b or c
+
+assert a or b
+assert c
+
+assert a or b
+
+assert a