Merge pull request #163 from bcaller/curry

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

Author: KevinHock

pyt/cfg/stmt_visitor.py

@@ -619,7 +619,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)
 
@@ -638,15 +639,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'])