feat: more ast helpers: __str__ and find_body (#136)

* feat: implement __str__

* feat: add find_body helper

* docs: describe str and find_body

* refactor: try to explain is_equivalent better

Also formatting.

Author: ojeytonwilliams

docs/python.md

@@ -122,6 +122,25 @@ assert add(a, b) == 300
 In contrast to the regex-based helpers, these helpers need to be run in Python, not
 JavaScript
 
+### Formatting output
+
+`str` returns a string that would parse to the same AST as the node. For example:
+
+```python
+function_str = """
+def foo():
+  # will not be in the output
+  x = 1
+
+"""
+output_str = """
+def foo():
+    x = 1"""
+str(Node(function_str)) == output_str # True
+```
+
+The output and source string compile to the same AST, but the output is indented with 4 spaces. Comments and trailing whitespace are removed.
+
 ### Finding nodes
 
 `find_` functions search the current scope and return one of the following:
@@ -147,6 +166,15 @@ When the variable is out of scope, `find_variable` returns an `None` node (i.e.
 Node('def foo():\n  x = "1"').find_variable("x") == Node() # True
 ```
 
+#### `find_body`
+
+```python
+func_str = """
+def foo():
+    x = 1"""
+Node(func_str).find_function("foo").find_body().is_equivalent("x = 1") # True
+```
+
 #### `find_class`
 
 ```python

python/py_helpers.py

@@ -41,6 +41,11 @@ def __repr__(self):
             return "Node:\nNone"
         return "Node:\n" + ast.dump(self.tree, indent=2)
 
+    def __str__(self):
+        if self.tree == None:
+            return "# no ast"
+        return ast.unparse(self.tree)
+
     def _has_body(self):
         return bool(getattr(self.tree, "body", False))
 
@@ -57,6 +62,13 @@ def find_function(self, func):
                     return Node(node)
         return Node()
 
+    def find_body(self):
+        if not isinstance(self.tree, ast.AST):
+            return Node()
+        if not hasattr(self.tree, "body"):
+            return Node()
+        return Node(ast.Module(self.tree.body, []))
+
     # "has" functions return a boolean indicating whether whatever is being
     # searched for exists. In this case, it returns True if the variable exists.
 
@@ -99,12 +111,18 @@ def value_is_call(self, name):
             return call.func.id == name
         return False
 
-    # Takes an string and checks if is equivalent to the node's AST. This
-    # is a loose comparison that tries to find out if the code is essentially
-    # the same. For example, the string "True" is not represented by the same
-    # AST as the test in "if True:" (the string could be wrapped in Module,
-    # Interactive or Expression, depending on the parse mode and the test is
-    # just a Constant), but they are equivalent.
+    # Loosely compares the code in target_str with the code represented by the
+    # Node's AST. If the two codes are semantically equivalent (i.e. the same if
+    # you ignore formatting and context) then this returns True, otherwise
+    # False.
+    #
+    # Ignoring context means that the following comparison is True despite the
+    # fact that the AST of `cond_node` is `Constant(value=True)` and `True`
+    # compiles to `Module(body=[Expr(value=Constant(value=True))],
+    # type_ignores=[])`:
+    #
+    # node = Node("if True:\n  pass") cond_node =
+    # node.find_ifs()[0].find_conditions()[0] cond_node.is_equivalent("True")
 
     def is_equivalent(self, target_str):
         # Setting the tree to None is used to represent missing elements. Such
@@ -131,9 +149,7 @@ def find_ifs(self):
         return self._find_all(ast.If)
 
     def _find_all(self, ast_type):
-        return [
-            Node(node) for node in self.tree.body if isinstance(node, ast_type)
-        ]
+        return [Node(node) for node in self.tree.body if isinstance(node, ast_type)]
 
     def find_conditions(self):
         def _find_conditions(tree):
@@ -162,6 +178,4 @@ def _find_if_bodies(tree):
 
             return [tree.body] + [tree.orelse]
 
-        return [
-            Node(ast.Module(body, [])) for body in _find_if_bodies(self.tree)
-        ]
+        return [Node(ast.Module(body, [])) for body in _find_if_bodies(self.tree)]

python/py_helpers.test.py

@@ -223,6 +223,37 @@ def test_not_has_function(self):
 
         self.assertFalse(node.has_function("bar"))
 
+    def test_find_body(self):
+        func_str = """def foo():
+  x = 1
+  print(x)
+"""
+        node = Node(func_str)
+
+        self.assertTrue(
+            node.find_function("foo").find_body().is_equivalent("x = 1\nprint(x)")
+        )
+        self.assertEqual("x = 1\nprint(x)", str(node.find_function("foo").find_body()))
+
+    def test_find_body_with_class(self):
+        class_str = """
+class Foo:
+  def __init__(self):
+    self.x = 1
+"""
+        node = Node(class_str)
+
+        self.assertTrue(
+            node.find_class("Foo")
+            .find_body()
+            .is_equivalent("def __init__(self):\n    self.x = 1")
+        )
+
+    def test_find_body_without_body(self):
+        node = Node("x = 1")
+
+        self.assertEqual(node.find_variable("x").find_body(), Node())
+
 
 class TestEquivalenceHelpers(unittest.TestCase):
     def test_is_equivalent(self):
@@ -478,6 +509,41 @@ def test_len(self):
 
         self.assertEqual(len(node.find_ifs()), 2)
 
+    def test_str(self):
+        func_str = """def foo():
+  pass
+"""
+        # Note: the indentation and whitespace is not preserved.
+        expected = """def foo():
+    pass"""
+
+        self.assertEqual(expected, str(Node(func_str)))
+
+    def test_none_str(self):
+        self.assertEqual("# no ast", str(Node()))
+
+    def test_str_with_comments(self):
+        func_str = """def foo():
+  # comment
+  pass
+
+
+"""
+        # Note: comments are discarded
+        expected = """def foo():
+    pass"""
+
+        self.assertEqual(expected, str(Node(func_str)))
+
+
+    def test_repr(self):
+        func_str = """def foo():
+  pass
+"""
+        node = Node(func_str)
+
+        self.assertEqual(repr(node), "Node:\n" + ast.dump(node.tree, indent=2))
+
 
 if __name__ == "__main__":
     unittest.main()