Repository: iris-hep/qastle

Issue:
`OrderBy` node type
Just add the simplest sorting operator in LINQ (<https://docs.microsoft.com/en-us/dotnet/api/system.linq.enumerable.orderby?view=net-5.0>). Here it would look like `(OrderBy <source> <keySelector>)`.

Comments:

Patch:
diff --git a/README.md b/README.md
index f6c520f..e97480f 100644
--- a/README.md
+++ b/README.md
@@ -76,6 +76,8 @@ All defined s-expressions are listed here, though this specification will be exp
   - Min: `(Min <source>)`
   - Sum: `(Sum <source>)`
   - Zip: `(Zip <source>)`
+  - OrderBy: `(OrderBy <source> <key_selector>)`
+    - `key_selector` must be a `lambda` with one argument
 
 
 ## Example
diff --git a/qastle/linq_util.py b/qastle/linq_util.py
index d9efe95..be3e88e 100644
--- a/qastle/linq_util.py
+++ b/qastle/linq_util.py
@@ -68,6 +68,13 @@ class Zip(ast.AST):
         self.source = source
 
 
+class OrderBy(ast.AST):
+    def __init__(self, source, key_selector):
+        self._fields = ['source', 'key_selector']
+        self.source = source
+        self.key_selector = key_selector
+
+
 linq_operator_names = ('Where',
                        'Select',
                        'SelectMany',
@@ -77,7 +84,8 @@ linq_operator_names = ('Where',
                        'Max',
                        'Min',
                        'Sum',
-                       'Zip')
+                       'Zip',
+                       'OrderBy')
 
 
 class InsertLINQNodesTransformer(ast.NodeTransformer):
@@ -161,6 +169,15 @@ class InsertLINQNodesTransformer(ast.NodeTransformer):
             if len(args) != 0:
                 raise SyntaxError('Zip() call must have zero arguments')
             return Zip(source=self.visit(source))
+        elif function_name == 'OrderBy':
+            if len(args) != 1:
+                raise SyntaxError('OrderBy() call must have exactly one argument')
+            if isinstance(args[0], ast.Str):
+                args[0] = unwrap_ast(ast.parse(args[0].s))
+            if not isinstance(args[0], ast.Lambda):
+                raise SyntaxError('OrderBy() call argument must be a lambda')
+            return OrderBy(source=self.visit(source),
+                           key_selector=self.visit(args[0]))
         else:
             raise NameError('Unhandled LINQ operator: ' + function_name)
 
diff --git a/qastle/transform.py b/qastle/transform.py
index 4017851..c0b45b7 100644
--- a/qastle/transform.py
+++ b/qastle/transform.py
@@ -1,4 +1,5 @@
-from .linq_util import Where, Select, SelectMany, First, Aggregate, Count, Max, Min, Sum, Zip
+from .linq_util import (Where, Select, SelectMany, First, Aggregate, Count, Max,
+                        Min, Sum, Zip, OrderBy)
 from .ast_util import wrap_ast, unwrap_ast
 
 import lark
@@ -203,6 +204,11 @@ class PythonASTToTextASTTransformer(ast.NodeVisitor):
     def visit_Zip(self, node):
         return self.make_composite_node_string('Zip', self.visit(node.source))
 
+    def visit_OrderBy(self, node):
+        return self.make_composite_node_string('OrderBy',
+                                               self.visit(node.source),
+                                               self.visit(node.key_selector))
+
     def generic_visit(self, node):
         raise SyntaxError('Unsupported node type: ' + str(type(node)))
 
@@ -438,5 +444,16 @@ class TextASTToPythonASTTransformer(lark.Transformer):
                 raise SyntaxError('Zip node must have one field; found ' + str(len(fields)))
             return Zip(source=fields[0])
 
+        elif node_type == 'OrderBy':
+            if len(fields) != 2:
+                raise SyntaxError('OrderBy node must have two fields; found ' + str(len(fields)))
+            if not isinstance(fields[1], ast.Lambda):
+                raise SyntaxError('OrderBy key selector must be a lambda; found '
+                                  + str(type(fields[1])))
+            if len(fields[1].args.args) != 1:
+                raise SyntaxError('OrderBy key selector must have exactly one argument; found '
+                                  + str(len(fields[1].args.args)))
+            return OrderBy(source=fields[0], key_selector=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 096575d..3404ce7 100644
--- a/tests/test_ast_language.py
+++ b/tests/test_ast_language.py
@@ -217,3 +217,10 @@ def test_Zip():
     first_node = Zip(source=unwrap_ast(ast.parse('data_source')))
     assert_equivalent_python_ast_and_text_ast(wrap_ast(first_node),
                                               '(Zip data_source)')
+
+
+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 data_source (lambda (list e) e))')
diff --git a/tests/test_linq_util.py b/tests/test_linq_util.py
index 93b1d07..d579eef 100644
--- a/tests/test_linq_util.py
+++ b/tests/test_linq_util.py
@@ -256,3 +256,26 @@ def test_zip_composite():
 def test_zip_bad():
     with pytest.raises(SyntaxError):
         insert_linq_nodes(ast.parse('the_source.Zip(None)'))
+
+
+def test_orderby():
+    initial_ast = ast.parse("the_source.OrderBy('lambda row: row')")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(OrderBy(source=unwrap_ast(ast.parse('the_source')),
+                                    key_selector=unwrap_ast(ast.parse('lambda row: row'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_orderby_composite():
+    initial_ast = ast.parse("the_source.First().OrderBy('lambda row: row')")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(OrderBy(source=First(source=unwrap_ast(ast.parse('the_source'))),
+                                    key_selector=unwrap_ast(ast.parse('lambda row: row'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_orderby_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.OrderBy()'))
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.OrderBy(None)'))