Repository: iris-hep/qastle

Issue:
`CrossJoin` node type
Should look like `(CrossJoin <first> <second>)`, where `first` and `second` are the sequences used to form a cartesian product.

Comments:

Patch:
diff --git a/README.md b/README.md
index e97480f..cbf062f 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,7 @@ All defined s-expressions are listed here, though this specification will be exp
   - Zip: `(Zip <source>)`
   - OrderBy: `(OrderBy <source> <key_selector>)`
     - `key_selector` must be a `lambda` with one argument
+  - CrossJoin: `(CrossJoin <first> <second>)`
 
 
 ## Example
diff --git a/qastle/linq_util.py b/qastle/linq_util.py
index be3e88e..9fb9251 100644
--- a/qastle/linq_util.py
+++ b/qastle/linq_util.py
@@ -75,6 +75,13 @@ class OrderBy(ast.AST):
         self.key_selector = key_selector
 
 
+class CrossJoin(ast.AST):
+    def __init__(self, first, second):
+        self._fields = ['first', 'second']
+        self.first = first
+        self.second = second
+
+
 linq_operator_names = ('Where',
                        'Select',
                        'SelectMany',
@@ -85,7 +92,8 @@ linq_operator_names = ('Where',
                        'Min',
                        'Sum',
                        'Zip',
-                       'OrderBy')
+                       'OrderBy',
+                       'CrossJoin')
 
 
 class InsertLINQNodesTransformer(ast.NodeTransformer):
@@ -178,6 +186,11 @@ class InsertLINQNodesTransformer(ast.NodeTransformer):
                 raise SyntaxError('OrderBy() call argument must be a lambda')
             return OrderBy(source=self.visit(source),
                            key_selector=self.visit(args[0]))
+        elif function_name == 'CrossJoin':
+            if len(args) != 1:
+                raise SyntaxError('CrossJoin() call must have exactly one argument')
+            return CrossJoin(first=self.visit(source),
+                             second=self.visit(args[0]))
         else:
             raise NameError('Unhandled LINQ operator: ' + function_name)
 
diff --git a/qastle/transform.py b/qastle/transform.py
index c0b45b7..fe9a820 100644
--- a/qastle/transform.py
+++ b/qastle/transform.py
@@ -1,5 +1,5 @@
 from .linq_util import (Where, Select, SelectMany, First, Aggregate, Count, Max,
-                        Min, Sum, Zip, OrderBy)
+                        Min, Sum, Zip, OrderBy, CrossJoin)
 from .ast_util import wrap_ast, unwrap_ast
 
 import lark
@@ -209,6 +209,11 @@ class PythonASTToTextASTTransformer(ast.NodeVisitor):
                                                self.visit(node.source),
                                                self.visit(node.key_selector))
 
+    def visit_CrossJoin(self, node):
+        return self.make_composite_node_string('CrossJoin',
+                                               self.visit(node.first),
+                                               self.visit(node.second))
+
     def generic_visit(self, node):
         raise SyntaxError('Unsupported node type: ' + str(type(node)))
 
@@ -455,5 +460,10 @@ class TextASTToPythonASTTransformer(lark.Transformer):
                                   + str(len(fields[1].args.args)))
             return OrderBy(source=fields[0], key_selector=fields[1])
 
+        elif node_type == 'CrossJoin':
+            if len(fields) != 2:
+                raise SyntaxError('CrossJoin node must have two fields; found ' + str(len(fields)))
+            return CrossJoin(first=fields[0], second=fields[1])
+
         else:
             raise SyntaxError('Unknown composite node type: ' + node_type)


Test patch:
diff --git a/tests/test_ast_language.py b/tests/test_ast_language.py
index 3404ce7..68aa6f8 100644
--- a/tests/test_ast_language.py
+++ b/tests/test_ast_language.py
@@ -220,7 +220,13 @@ def test_Zip():
 
 
 def test_OrderBy():
-    select_node = OrderBy(source=unwrap_ast(ast.parse('data_source')),
-                          key_selector=unwrap_ast(ast.parse('lambda e: e')))
-    assert_equivalent_python_ast_and_text_ast(wrap_ast(select_node),
+    orderby_node = OrderBy(source=unwrap_ast(ast.parse('data_source')),
+                           key_selector=unwrap_ast(ast.parse('lambda e: e')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(orderby_node),
                                               '(OrderBy data_source (lambda (list e) e))')
+
+
+def test_CrossJoin():
+    crossjoin_node = CrossJoin(first=unwrap_ast(ast.parse('left')),
+                               second=unwrap_ast(ast.parse('right')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(crossjoin_node), '(CrossJoin left right)')
diff --git a/tests/test_linq_util.py b/tests/test_linq_util.py
index d579eef..a4576d2 100644
--- a/tests/test_linq_util.py
+++ b/tests/test_linq_util.py
@@ -279,3 +279,24 @@ def test_orderby_bad():
         insert_linq_nodes(ast.parse('the_source.OrderBy()'))
     with pytest.raises(SyntaxError):
         insert_linq_nodes(ast.parse('the_source.OrderBy(None)'))
+
+
+def test_crossjoin():
+    initial_ast = ast.parse("left.CrossJoin(right)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(CrossJoin(first=unwrap_ast(ast.parse('left')),
+                                      second=unwrap_ast(ast.parse('right'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_crossjoin_composite():
+    initial_ast = ast.parse("left.First().CrossJoin(right)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(CrossJoin(first=First(source=unwrap_ast(ast.parse('left'))),
+                                      second=unwrap_ast(ast.parse('right'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_crossjoin_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('left.CrossJoin()'))