Repository: iris-hep/qastle

Issue:
`ElementAt` node type
See <https://docs.microsoft.com/en-us/dotnet/api/system.linq.enumerable.elementat?view=net-5.0>. For parity with `First` and `Last`.

Comments:

Patch:
diff --git a/README.md b/README.md
index 27155e5..55ff7d2 100644
--- a/README.md
+++ b/README.md
@@ -70,6 +70,8 @@ All defined s-expressions are listed here, though this specification will be exp
     - `selector` must be a `lambda` with one argument
   - First: `(First <source>)`
   - Last: `(Last <source>)`
+  - ElementAt: `(ElementAt <source> <index>)`
+    - `index` must be an integer
   - Aggregate: `(Aggregate <source> <seed> <func>)`
     - `func` must be a `lambda` with two arguments
   - Count: `(Count <source>)`
diff --git a/qastle/linq_util.py b/qastle/linq_util.py
index df6eb79..a4922ec 100644
--- a/qastle/linq_util.py
+++ b/qastle/linq_util.py
@@ -23,6 +23,10 @@ class Last(ast.AST):
     _fields = ['source']
 
 
+class ElementAt(ast.AST):
+    _fields = ['source', 'index']
+
+
 class Aggregate(ast.AST):
     _fields = ['source', 'seed', 'func']
 
@@ -60,6 +64,7 @@ linq_operator_names = ('Where',
                        'SelectMany',
                        'First',
                        'Last',
+                       'ElementAt',
                        'Aggregate',
                        'Count',
                        'Max',
@@ -124,6 +129,10 @@ class InsertLINQNodesTransformer(ast.NodeTransformer):
             if len(args) != 0:
                 raise SyntaxError('Last() call must have zero arguments')
             return Last(source=self.visit(source))
+        elif function_name == 'ElementAt':
+            if len(args) != 1:
+                raise SyntaxError('ElementAt() call must have exactly one argument')
+            return ElementAt(source=self.visit(source), index=self.visit(args[0]))
         elif function_name == 'Aggregate':
             if len(args) != 2:
                 raise SyntaxError('Aggregate() call must have exactly two arguments; found'
diff --git a/qastle/transform.py b/qastle/transform.py
index e13b16d..da7dab8 100644
--- a/qastle/transform.py
+++ b/qastle/transform.py
@@ -1,4 +1,4 @@
-from .linq_util import (Where, Select, SelectMany, First, Last, Aggregate, Count, Max,
+from .linq_util import (Where, Select, SelectMany, First, Last, ElementAt, Aggregate, Count, Max,
                         Min, Sum, Zip, OrderBy, Choose)
 from .ast_util import wrap_ast, unwrap_ast
 
@@ -186,6 +186,11 @@ class PythonASTToTextASTTransformer(ast.NodeVisitor):
     def visit_Last(self, node):
         return self.make_composite_node_string('Last', self.visit(node.source))
 
+    def visit_ElementAt(self, node):
+        return self.make_composite_node_string('ElementAt',
+                                               self.visit(node.source),
+                                               self.visit(node.index))
+
     def visit_Aggregate(self, node):
         return self.make_composite_node_string('Aggregate',
                                                self.visit(node.source),
@@ -421,6 +426,11 @@ class TextASTToPythonASTTransformer(lark.Transformer):
                 raise SyntaxError('Last node must have one field; found ' + str(len(fields)))
             return Last(source=fields[0])
 
+        elif node_type == 'ElementAt':
+            if len(fields) != 2:
+                raise SyntaxError('ElementAt node must have two fields; found ' + str(len(fields)))
+            return ElementAt(source=fields[0], index=fields[1])
+
         elif node_type == 'Aggregate':
             if len(fields) != 3:
                 raise SyntaxError('Aggregate node must have three fields; found '


Test patch:
diff --git a/tests/test_ast_language.py b/tests/test_ast_language.py
index b83cb8d..a9fdec9 100644
--- a/tests/test_ast_language.py
+++ b/tests/test_ast_language.py
@@ -180,6 +180,19 @@ def test_First():
                                               '(First data_source)')
 
 
+def test_Last():
+    first_node = Last(source=unwrap_ast(ast.parse('data_source')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(first_node),
+                                              '(Last data_source)')
+
+
+def test_ElementAt():
+    first_node = ElementAt(source=unwrap_ast(ast.parse('data_source')),
+                           index=unwrap_ast(ast.parse('2')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(first_node),
+                                              '(ElementAt data_source 2)')
+
+
 def test_Aggregate():
     aggregate_node = Aggregate(source=unwrap_ast(ast.parse('data_source')),
                                seed=unwrap_ast(ast.parse('0')),
@@ -231,9 +244,3 @@ def test_Choose():
                          n=unwrap_ast(ast.parse('2')))
     assert_equivalent_python_ast_and_text_ast(wrap_ast(choose_node),
                                               '(Choose data_source 2)')
-
-
-def test_Last():
-    first_node = Last(source=unwrap_ast(ast.parse('data_source')))
-    assert_equivalent_python_ast_and_text_ast(wrap_ast(first_node),
-                                              '(Last data_source)')
diff --git a/tests/test_linq_util.py b/tests/test_linq_util.py
index cfc7e9c..7ce2580 100644
--- a/tests/test_linq_util.py
+++ b/tests/test_linq_util.py
@@ -150,6 +150,46 @@ def test_first_bad():
         insert_linq_nodes(ast.parse('the_source.First(None)'))
 
 
+def test_last():
+    initial_ast = ast.parse("the_source.Last()")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Last(source=unwrap_ast(ast.parse('the_source'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_last_composite():
+    initial_ast = ast.parse("the_source.First().Last()")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Last(source=First(source=unwrap_ast(ast.parse('the_source')))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_last_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.Last(None)'))
+
+
+def test_elementat():
+    initial_ast = ast.parse("the_source.ElementAt(2)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(ElementAt(source=unwrap_ast(ast.parse('the_source')),
+                                      index=unwrap_ast(ast.parse('2'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_elementat_composite():
+    initial_ast = ast.parse("the_source.First().ElementAt(2)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(ElementAt(source=First(source=unwrap_ast(ast.parse('the_source'))),
+                                      index=unwrap_ast(ast.parse('2'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_elementat_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.ElementAt()'))
+
+
 def test_aggregate():
     initial_ast = ast.parse("the_source.Aggregate(0, 'lambda total, next: total + next')")
     final_ast = insert_linq_nodes(initial_ast)
@@ -316,22 +356,3 @@ def test_choose_composite():
 def test_choose_bad():
     with pytest.raises(SyntaxError):
         insert_linq_nodes(ast.parse('the_source.Choose()'))
-
-
-def test_last():
-    initial_ast = ast.parse("the_source.Last()")
-    final_ast = insert_linq_nodes(initial_ast)
-    expected_ast = wrap_ast(Last(source=unwrap_ast(ast.parse('the_source'))))
-    assert_ast_nodes_are_equal(final_ast, expected_ast)
-
-
-def test_last_composite():
-    initial_ast = ast.parse("the_source.Last().Last()")
-    final_ast = insert_linq_nodes(initial_ast)
-    expected_ast = wrap_ast(Last(source=Last(source=unwrap_ast(ast.parse('the_source')))))
-    assert_ast_nodes_are_equal(final_ast, expected_ast)
-
-
-def test_last_bad():
-    with pytest.raises(SyntaxError):
-        insert_linq_nodes(ast.parse('the_source.Last(None)'))