Fix VarsVisitor RuntimeError on code like f(g(a)(b)(c))

pyt/cfg/stmt_visitor.py

@@ -615,7 +615,8 @@ def add_blackbox_or_builtin_call(self, node, blackbox):
         rhs_vars = list()
         last_return_value_of_nested_call = None
 
-        for arg in itertools.chain(node.args, node.keywords):
+        for arg_node in itertools.chain(node.args, node.keywords):
+            arg = arg_node.value if isinstance(arg_node, ast.keyword) else arg_node
             if isinstance(arg, ast.Call):
                 return_value_of_nested_call = self.visit(arg)
 
@@ -634,15 +635,18 @@ def add_blackbox_or_builtin_call(self, node, blackbox):
                     call_node.inner_most_call = return_value_of_nested_call
                 last_return_value_of_nested_call = return_value_of_nested_call
 
-                visual_args.append(return_value_of_nested_call.left_hand_side)
+                if isinstance(arg_node, ast.keyword) and arg_node.arg is not None:
+                    visual_args.append(arg_node.arg + '=' + return_value_of_nested_call.left_hand_side)
+                else:
+                    visual_args.append(return_value_of_nested_call.left_hand_side)
                 rhs_vars.append(return_value_of_nested_call.left_hand_side)
             else:
                 label = LabelVisitor()
-                label.visit(arg)
+                label.visit(arg_node)
                 visual_args.append(label.result)
 
                 vv = VarsVisitor()
-                vv.visit(arg)
+                vv.visit(arg_node)
                 rhs_vars.extend(vv.result)
         if last_return_value_of_nested_call:
             # connect other_inner to outer in e.g.

pyt/helper_visitors/vars_visitor.py

@@ -84,7 +84,8 @@ def visit_Call(self, node):
         # This will not visit Flask in Flask(__name__) but it will visit request in `request.args.get()
         if not isinstance(node.func, ast.Name):
             self.visit(node.func)
-        for arg in itertools.chain(node.args, node.keywords):
+        for arg_node in itertools.chain(node.args, node.keywords):
+            arg = arg_node.value if isinstance(arg_node, ast.keyword) else arg_node
             if isinstance(arg, ast.Call):
                 if isinstance(arg.func, ast.Name):
                     # We can't just visit because we need to add 'ret_'
@@ -95,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

@@ -45,9 +45,20 @@ def test_call5(self):
         self.assertEqual(vars.result, ['resp', 'ret_replace'])
 
     def test_call6(self):
+        vars = self.perform_vars_on_expression("resp = f(kw=g(a, b))")
+        self.assertEqual(vars.result, ['resp', 'ret_g'])
+
+    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'])