Don't crash on pathological case of f(g(a)(b)(c))

It's rare, but a curried function call can appear within a function
call. It was raising a RuntimeError in VarsVisitor.

We don't build a cfg properly for curried functions which is fine for
now, but we don't need to crash.

At least there is now defined behaviour.

Author: bcaller

pyt/helper_visitors/vars_visitor.py

@@ -96,12 +96,32 @@ def visit_Call(self, node):
                     # func.value.id is html
                     # We want replace
                     self.result.append('ret_' + arg.func.attr)
+                elif isinstance(arg.func, ast.Call):
+                    self.visit_curried_call_inside_call_args(arg)
                 else:
-                    # Deal with it when we have code that triggers it.
-                    raise
+                    raise Exception('Cannot visit vars of ' + ast.dump(arg))
             else:
                 self.visit(arg)
 
+    def visit_curried_call_inside_call_args(self, inner_call):
+        # Curried functions aren't supported really, but we now at least have a defined behaviour.
+        # In f(g(a)(b)(c)), inner_call is the Call node with argument c
+        # Try to get the name of curried function g
+        curried_func = inner_call.func.func
+        while isinstance(curried_func, ast.Call):
+            curried_func = curried_func.func
+        if isinstance(curried_func, ast.Name):
+            self.result.append('ret_' + curried_func.id)
+        elif isinstance(curried_func, ast.Attribute):
+            self.result.append('ret_' + curried_func.attr)
+
+        # Visit all arguments except a (ignore the curried function g)
+        not_curried = inner_call
+        while not_curried.func is not curried_func:
+            for arg in itertools.chain(not_curried.args, not_curried.keywords):
+                self.visit(arg.value if isinstance(arg, ast.keyword) else arg)
+            not_curried = not_curried.func
+
     def visit_Attribute(self, node):
         if not isinstance(node.value, ast.Name):
             self.visit(node.value)

tests/helper_visitors/vars_visitor_test.py

@@ -52,6 +52,13 @@ def test_call7(self):
         vars = self.perform_vars_on_expression("resp = make_response(html.replace.bar('{{ param }}', param))")
         self.assertEqual(vars.result, ['resp', 'ret_bar'])
 
+    def test_curried_function(self):
+        # Curried functions aren't supported really, but we now at least have a defined behaviour.
+        vars = self.perform_vars_on_expression('f(g.h(a)(b))')
+        self.assertCountEqual(vars.result, ['ret_h', 'b'])
+        vars = self.perform_vars_on_expression('f(g(a)(b)(c)(d, e=f))')
+        self.assertCountEqual(vars.result, ['ret_g', 'b', 'c', 'd', 'f'])
+
     def test_keyword_vararg(self):
         vars = self.perform_vars_on_expression('print(arg = x)')
         self.assertEqual(vars.result, ['x'])