Repository: iris-hep/qastle

Issue:
`Concat` node type
As in <https://docs.microsoft.com/en-us/dotnet/api/system.linq.enumerable.concat?view=net-5.0>.

Comments:

Patch:
diff --git a/README.md b/README.md
index 4b9d7d5..2c040fb 100644
--- a/README.md
+++ b/README.md
@@ -82,6 +82,7 @@ All defined s-expressions are listed here, though this specification will be exp
     - `predicate` must be a `lambda` with one argument
   - Any: `(Any <source> <predicate>)`
     - `predicate` must be a `lambda` with one argument
+  - Concat: `(Concat <first> <second>)`
   - Zip: `(Zip <source>)`
   - OrderBy: `(OrderBy <source> <key_selector>)`
     - `key_selector` must be a `lambda` with one argument
diff --git a/qastle/linq_util.py b/qastle/linq_util.py
index f04fcbe..bb0e2ff 100644
--- a/qastle/linq_util.py
+++ b/qastle/linq_util.py
@@ -55,6 +55,10 @@ class Any(ast.AST):
     _fields = ['source', 'predicate']
 
 
+class Concat(ast.AST):
+    _fields = ['first', 'second']
+
+
 class Zip(ast.AST):
     _fields = ['source']
 
@@ -84,6 +88,7 @@ linq_operator_names = ('Where',
                        'Sum',
                        'All',
                        'Any',
+                       'Concat',
                        'Zip',
                        'OrderBy',
                        'OrderByDescending',
@@ -193,6 +198,10 @@ class InsertLINQNodesTransformer(ast.NodeTransformer):
                 raise SyntaxError('Any() call argument must be a lambda')
             return Any(source=self.visit(source),
                        predicate=self.visit(args[0]))
+        elif function_name == 'Concat':
+            if len(args) != 1:
+                raise SyntaxError('Concat() call must have exactly one argument')
+            return Concat(first=self.visit(source), second=self.visit(args[0]))
         elif function_name == 'Zip':
             if len(args) != 0:
                 raise SyntaxError('Zip() call must have zero arguments')
diff --git a/qastle/transform.py b/qastle/transform.py
index f109686..f1c2629 100644
--- a/qastle/transform.py
+++ b/qastle/transform.py
@@ -1,5 +1,5 @@
 from .linq_util import (Where, Select, SelectMany, First, Last, ElementAt, Aggregate, Count, Max,
-                        Min, Sum, All, Any, Zip, OrderBy, OrderByDescending, Choose)
+                        Min, Sum, All, Any, Concat, Zip, OrderBy, OrderByDescending, Choose)
 from .ast_util import wrap_ast, unwrap_ast
 
 import lark
@@ -219,6 +219,11 @@ class PythonASTToTextASTTransformer(ast.NodeVisitor):
                                                self.visit(node.source),
                                                self.visit(node.predicate))
 
+    def visit_Concat(self, node):
+        return self.make_composite_node_string('Concat',
+                                               self.visit(node.first),
+                                               self.visit(node.second))
+
     def visit_Zip(self, node):
         return self.make_composite_node_string('Zip', self.visit(node.source))
 
@@ -499,6 +504,11 @@ class TextASTToPythonASTTransformer(lark.Transformer):
                                   + str(len(fields[1].args.args)))
             return Any(source=fields[0], predicate=fields[1])
 
+        elif node_type == 'Concat':
+            if len(fields) != 2:
+                raise SyntaxError('Concat node must have two fields; found ' + str(len(fields)))
+            return Concat(first=fields[0], second=fields[1])
+
         elif node_type == 'Zip':
             if len(fields) != 1:
                 raise SyntaxError('Zip node must have one field; found ' + str(len(fields)))


Test patch:
diff --git a/tests/test_ast_language.py b/tests/test_ast_language.py
index 5c340e8..1597b55 100644
--- a/tests/test_ast_language.py
+++ b/tests/test_ast_language.py
@@ -240,6 +240,13 @@ def test_Any():
                                               '(Any data_source (lambda (list e) e))')
 
 
+def test_Concat():
+    first_node = Concat(first=unwrap_ast(ast.parse('sequence1')),
+                        second=unwrap_ast(ast.parse('sequence2')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(first_node),
+                                              '(Concat sequence1 sequence2)')
+
+
 def test_Zip():
     first_node = Zip(source=unwrap_ast(ast.parse('data_source')))
     assert_equivalent_python_ast_and_text_ast(wrap_ast(first_node),
diff --git a/tests/test_linq_util.py b/tests/test_linq_util.py
index 65b3893..cef83c9 100644
--- a/tests/test_linq_util.py
+++ b/tests/test_linq_util.py
@@ -341,6 +341,27 @@ def test_any_bad():
         insert_linq_nodes(ast.parse('the_source.Any(None)'))
 
 
+def test_concat():
+    initial_ast = ast.parse("sequence1.Concat(sequence2)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Concat(first=unwrap_ast(ast.parse('sequence1')),
+                                   second=unwrap_ast(ast.parse('sequence2'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_concat_composite():
+    initial_ast = ast.parse("the_source.First().Concat(sequence)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Concat(first=First(source=unwrap_ast(ast.parse('the_source'))),
+                                   second=unwrap_ast(ast.parse('sequence'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_concat_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.Concat()'))
+
+
 def test_zip():
     initial_ast = ast.parse("the_source.Zip()")
     final_ast = insert_linq_nodes(initial_ast)