Make get_call_names more resilient

Ben Caller · Ben Caller · commit 83e496f78a3d · 2018-07-19T15:44:36.000+01:00
Code where we call a method of something which isn't a variable name
e.g.

```
yesterday = (date.today() - timedelta(days=1)).strftime("%Y-%m-%d")
```

was causing pyt to crash.

Previously the else branch would only handle ast.Attribute, and crash on
everything else.

```
Traceback (most recent call last):
  File "/usr/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/pyt/pyt/__main__.py", line 141, in &lt;module&gt;
    main()
  File "/pyt/pyt/__main__.py", line 97, in main
    path
  File "/pyt/pyt/cfg/make_cfg.py", line 40, in make_cfg
    module_definitions
  File "/pyt/pyt/cfg/expr_visitor.py", line 60, in __init__
    self.init_cfg(node)
  File "/pyt/pyt/cfg/expr_visitor.py", line 67, in init_cfg
    module_statements = self.visit(node)
  File "/usr/lib/python3.6/ast.py", line 253, in visit
    return visitor(node)
  File "/pyt/pyt/cfg/stmt_visitor.py", line 58, in visit_Module
    return self.stmt_star_handler(node.body)
  File "/pyt/pyt/cfg/stmt_visitor.py", line 79, in stmt_star_handler
    node = self.visit(stmt)
  File "/usr/lib/python3.6/ast.py", line 253, in visit
    return visitor(node)
  File "/pyt/pyt/cfg/stmt_visitor.py", line 413, in visit_Assign
    return self.assignment_call_node(label.result, node)
  File "/pyt/pyt/cfg/stmt_visitor.py", line 437, in assignment_call_node
    call = self.visit(ast_node.value)
  File "/usr/lib/python3.6/ast.py", line 253, in visit
    return visitor(node)
  File "/pyt/pyt/cfg/expr_visitor.py", line 540, in visit_Call
    _id = get_call_names_as_string(node.func)
  File "/pyt/pyt/core/ast_helper.py", line 78, in get_call_names_as_string
    return _list_to_dotted_string(get_call_names(node))
  File "/pyt/pyt/core/ast_helper.py", line 68, in get_call_names
    return reversed(_get_call_names_helper(node, result))
  File "/pyt/pyt/core/ast_helper.py", line 62, in _get_call_names_helper
    return _get_call_names_helper(node.value, result)
  File "/pyt/pyt/core/ast_helper.py", line 61, in _get_call_names_helper
    result.append(node.attr)
AttributeError: 'BinOp' object has no attribute 'attr'
```
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')
