Repository: iris-hep/qastle

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

Discussion:

Patch:
diff --git a/README.md b/README.md
index 320cd55..c2a0ab1 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,10 @@ All defined s-expressions are listed here, though this specification will be exp
   - Max: `(Max <source>)`
   - Min: `(Min <source>)`
   - Sum: `(Sum <source>)`
+  - All: `(All <source> <predicate>)`
+    - `predicate` must be a `lambda` with one argument
+  - Any: `(Any <source> <predicate>)`
+    - `predicate` must be a `lambda` with one argument
   - 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 b8ee702..f04fcbe 100644
--- a/qastle/linq_util.py
+++ b/qastle/linq_util.py
@@ -47,6 +47,14 @@ class Sum(ast.AST):
     _fields = ['source']
 
 
+class All(ast.AST):
+    _fields = ['source', 'predicate']
+
+
+class Any(ast.AST):
+    _fields = ['source', 'predicate']
+
+
 class Zip(ast.AST):
     _fields = ['source']
 
@@ -74,6 +82,8 @@ linq_operator_names = ('Where',
                        'Max',
                        'Min',
                        'Sum',
+                       'All',
+                       'Any',
                        'Zip',
                        'OrderBy',
                        'OrderByDescending',
@@ -165,6 +175,24 @@ class InsertLINQNodesTransformer(ast.NodeTransformer):
             if len(args) != 0:
                 raise SyntaxError('Sum() call must have zero arguments')
             return Sum(source=self.visit(source))
+        elif function_name == 'All':
+            if len(args) != 1:
+                raise SyntaxError('All() 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('All() call argument must be a lambda')
+            return All(source=self.visit(source),
+                       predicate=self.visit(args[0]))
+        elif function_name == 'Any':
+            if len(args) != 1:
+                raise SyntaxError('Any() 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('Any() call argument must be a lambda')
+            return Any(source=self.visit(source),
+                       predicate=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 76a32d5..f109686 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, Zip, OrderBy, OrderByDescending, Choose)
+                        Min, Sum, All, Any, Zip, OrderBy, OrderByDescending, Choose)
 from .ast_util import wrap_ast, unwrap_ast
 
 import lark
@@ -209,6 +209,16 @@ class PythonASTToTextASTTransformer(ast.NodeVisitor):
     def visit_Sum(self, node):
         return self.make_composite_node_string('Sum', self.visit(node.source))
 
+    def visit_All(self, node):
+        return self.make_composite_node_string('All',
+                                               self.visit(node.source),
+                                               self.visit(node.predicate))
+
+    def visit_Any(self, node):
+        return self.make_composite_node_string('Any',
+                                               self.visit(node.source),
+                                               self.visit(node.predicate))
+
     def visit_Zip(self, node):
         return self.make_composite_node_string('Zip', self.visit(node.source))
 
@@ -467,6 +477,28 @@ class TextASTToPythonASTTransformer(lark.Transformer):
                 raise SyntaxError('Sum node must have one field; found ' + str(len(fields)))
             return Sum(source=fields[0])
 
+        elif node_type == 'All':
+            if len(fields) != 2:
+                raise SyntaxError('All node must have two fields; found ' + str(len(fields)))
+            if not isinstance(fields[1], ast.Lambda):
+                raise SyntaxError('All predicate must be a lambda; found '
+                                  + str(type(fields[1])))
+            if len(fields[1].args.args) != 1:
+                raise SyntaxError('All predicate must have exactly one argument; found '
+                                  + str(len(fields[1].args.args)))
+            return All(source=fields[0], predicate=fields[1])
+
+        elif node_type == 'Any':
+            if len(fields) != 2:
+                raise SyntaxError('Any node must have two fields; found ' + str(len(fields)))
+            if not isinstance(fields[1], ast.Lambda):
+                raise SyntaxError('Any predicate must be a lambda; found '
+                                  + str(type(fields[1])))
+            if len(fields[1].args.args) != 1:
+                raise SyntaxError('Any predicate must have exactly one argument; found '
+                                  + str(len(fields[1].args.args)))
+            return Any(source=fields[0], predicate=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 991d913..5c340e8 100644
--- a/tests/test_ast_language.py
+++ b/tests/test_ast_language.py
@@ -226,6 +226,20 @@ def test_Sum():
                                               '(Sum data_source)')
 
 
+def test_All():
+    select_node = All(source=unwrap_ast(ast.parse('data_source')),
+                      predicate=unwrap_ast(ast.parse('lambda e: e')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(select_node),
+                                              '(All data_source (lambda (list e) e))')
+
+
+def test_Any():
+    select_node = Any(source=unwrap_ast(ast.parse('data_source')),
+                      predicate=unwrap_ast(ast.parse('lambda e: e')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(select_node),
+                                              '(Any data_source (lambda (list e) e))')
+
+
 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 2cf83c4..65b3893 100644
--- a/tests/test_linq_util.py
+++ b/tests/test_linq_util.py
@@ -295,6 +295,52 @@ def test_sum_bad():
         insert_linq_nodes(ast.parse('the_source.Sum(None)'))
 
 
+def test_all():
+    initial_ast = ast.parse("the_source.All('lambda row: row')")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(All(source=unwrap_ast(ast.parse('the_source')),
+                                predicate=unwrap_ast(ast.parse('lambda row: row'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_all_composite():
+    initial_ast = ast.parse("the_source.First().All('lambda row: row')")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(All(source=First(source=unwrap_ast(ast.parse('the_source'))),
+                                predicate=unwrap_ast(ast.parse('lambda row: row'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_all_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.All()'))
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.All(None)'))
+
+
+def test_any():
+    initial_ast = ast.parse("the_source.Any('lambda row: row')")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Any(source=unwrap_ast(ast.parse('the_source')),
+                                predicate=unwrap_ast(ast.parse('lambda row: row'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_any_composite():
+    initial_ast = ast.parse("the_source.First().Any('lambda row: row')")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Any(source=First(source=unwrap_ast(ast.parse('the_source'))),
+                                predicate=unwrap_ast(ast.parse('lambda row: row'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_any_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.Any()'))
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.Any(None)'))
+
+
 def test_zip():
     initial_ast = ast.parse("the_source.Zip()")
     final_ast = insert_linq_nodes(initial_ast)