Merge pull request #84 from iris-hep/fix_83_fields-of-choose

Fix fields when used with `Choose`

Author: masonproffitt

func_adl_uproot/transformer.py

@@ -42,6 +42,7 @@
 class PythonSourceGeneratorTransformer(ast.NodeTransformer):
     def __init__(self):
         self._depth = None
+        self._tuple_depths = []
         self._id_scopes = {}
         self._projection_stack = []
 
@@ -316,9 +317,15 @@ def visit_Select(self, node):
             raise TypeError('Lambda function in Select() must have exactly one argument, found '
                             + len(node.selector.args.args))
         self.visit(node.source)
+        if self._depth in self._tuple_depths:
+            at_tuple = True
+            original_source_rep = self.get_rep(node.source)
+            node.source.rep = 'x'
+        else:
+            at_tuple = False
         self._depth += 1
         self._projection_stack.append(node.selector.args.args[0].arg)
-        if self._depth > 2:
+        if self._depth > 2 and not at_tuple:
             rep1, rep2 = self._projection_stack[-2], self._projection_stack[-1]
             lambda_node = ast.Lambda(args=ast.arguments(args=[ast.arg(arg=rep1),
                                                               ast.arg(arg=rep2)]),
@@ -329,12 +336,17 @@ def visit_Select(self, node):
         else:
             call_node = ast.Call(func=node.selector, args=[node.source])
             call_rep = self.get_rep(call_node)
-        node.rep = ('(lambda selection: ak.zip(selection,'
-                    + ' depth_limit=(None if len(selection) == 1 else ' + repr(self._depth) + '))'
-                    + ' if not isinstance(selection, ak.Array)'
-                    + ' else selection)(' + call_rep + ')')
+        select_rep = ('(lambda selection: ak.zip(selection,'
+                      + ' depth_limit=(None if len(selection) == 1 else ' + repr(self._depth)
+                      + '))' + ' if not isinstance(selection, ak.Array)'
+                      + ' else selection)(' + call_rep + ')')
         self._depth -= 1
         self._projection_stack.pop()
+        if at_tuple:
+            node.rep = ('ak.zip([' + select_rep + ' for x in ak.unzip(' + original_source_rep
+                        + ')])')
+        else:
+            node.rep = select_rep
         return node
 
     def visit_SelectMany(self, node):
@@ -419,6 +431,7 @@ def visit_Choose(self, node):
         self.visit(node.source)
         node.rep = ('ak.combinations(' + self.get_rep(node.source) + ', ' + self.get_rep(node.n)
                     + ', axis=' + repr(self._depth) + ')')
+        self._tuple_depths.append(self._depth + 1)
         return node
 
     def visit_OrderBy(self, node, ascending=True):

tests/test_executor.py

@@ -326,6 +326,25 @@ def test_ast_executor_select_of_choose():
     assert ast_executor(python_ast).tolist() == [[], [1, 2, 5], []]
 
 
+def test_ast_executor_choose_zipped_dict():
+    python_source = ("Select(EventDataset('tests/vectors_tree_file.root', 'tree'),"
+                     + "lambda row: {'ints': row.int_vector_branch}.Zip().Choose(2))")
+    python_ast = ast.parse(python_source)
+    assert ast_executor(python_ast).tolist() == [[],
+                                                 [({'ints': -1}, {'ints': 2}),
+                                                  ({'ints': -1}, {'ints': 3}),
+                                                  ({'ints': 2}, {'ints': 3})],
+                                                 []]
+
+
+def test_ast_executor_field_of_choose():
+    python_source = ("Select(EventDataset('tests/vectors_tree_file.root', 'tree'),"
+                     + "lambda row: {'int': row.int_vector_branch}.Zip().Choose(2)"
+                     + '.Select(lambda pair: pair.Select(lambda record: record.int)))')
+    python_ast = ast.parse(python_source)
+    assert ast_executor(python_ast).tolist() == [[], [(-1, 2), (-1, 3), (2, 3)], []]
+
+
 def test_ast_executor_tofourvector():
     python_source = ("Select(EventDataset('tests/four-vector_tree_file.root', 'tree'),"
                      + "lambda row: Zip({'pt': row.pt_vector_branch, 'eta': row.eta_vector_branch,"