Merge pull request #148 from bcaller/bc-get-call-names

KevinHock · web-flow · commit cbc17f665441 · 2018-07-19T18:11:32.000-07:00
Make get_call_names more resilient
diff --git a/pyt/core/ast_helper.py b/pyt/core/ast_helper.py
@@ -44,28 +44,23 @@ def generate_ast(path):
     raise IOError('Input needs to be a file. Path: ' + path)
 
 
-def _get_call_names_helper(node, result):
+def _get_call_names_helper(node):
     """Recursively finds all function names."""
     if isinstance(node, ast.Name):
         if node.id not in BLACK_LISTED_CALL_NAMES:
-            result.append(node.id)
-        return result
-    elif isinstance(node, ast.Call):
-        return result
+            yield node.id
     elif isinstance(node, ast.Subscript):
-        return _get_call_names_helper(node.value, result)
+        yield from _get_call_names_helper(node.value)
     elif isinstance(node, ast.Str):
-        result.append(node.s)
-        return result
-    else:
-        result.append(node.attr)
-        return _get_call_names_helper(node.value, result)
+        yield node.s
+    elif isinstance(node, ast.Attribute):
+        yield node.attr
+        yield from _get_call_names_helper(node.value)
 
 
 def get_call_names(node):
     """Get a list of call names."""
-    result = list()
-    return reversed(_get_call_names_helper(node, result))
+    return reversed(list(_get_call_names_helper(node)))
 
 
 def _list_to_dotted_string(list_of_components):
diff --git a/tests/cfg/import_test.py b/tests/cfg/import_test.py
@@ -733,3 +733,19 @@ def test_get_call_names_multi(self):
         result = get_call_names_as_string(call.func)
 
         self.assertEqual(result, 'abc.defg.hi')
+
+    def test_get_call_names_with_binop(self):
+        m = ast.parse('(date.today() - timedelta(days=1)).strftime("%Y-%m-%d")')
+        call = m.body[0].value
+
+        result = get_call_names_as_string(call.func)
+
+        self.assertEqual(result, 'strftime')
+
+    def test_get_call_names_with_comprehension(self):
+        m = ast.parse('{a for a in b()}.union(c)')
+        call = m.body[0].value
+
+        result = get_call_names_as_string(call.func)
+
+        self.assertEqual(result, 'union')
