Merge pull request #190 from iris-hep:big/range

Make sure Range isn't Replaced

Author: gordonwatts

func_adl/functions.py

@@ -1,13 +1,12 @@
 import ast
-from typing import Union, cast
+from typing import Union
 
 from .object_stream import ObjectStream
-from .util_ast import as_ast, function_call
 
 
 def Range(
     lower_bound: Union[str, int, ast.AST], upper_bound: Union[str, int, ast.AST]
-) -> ObjectStream:
+) -> ObjectStream[int]:
     r"""
     Given the lower and upper bound return an object with the range of numbers, similar to python
     range
@@ -20,11 +19,5 @@ def Range(
     Return:
 
         A new ObjectStream that contains the range of numbers
-
     """
-
-    return ObjectStream(
-        function_call(
-            "Range", [cast(ast.AST, as_ast(lower_bound)), cast(ast.AST, as_ast(upper_bound))]
-        )
-    )
+    ...

tests/test_util_ast.py

@@ -1091,3 +1091,17 @@ def __init__(self, n):
         check_ast(a)
 
     assert "my_type" in str(e)
+
+
+def test_Select_inline():
+    "Make sure nothing funny happens to Select"
+    r = parse_as_ast(lambda jets: jets.Select(lambda j: j.pt()))
+    assert "function_call" not in ast.unparse(r)
+
+
+def test_Range_inline():
+    "Make sure nothing funny happens to Range"
+    from func_adl import Range
+
+    r = parse_as_ast(lambda jets: Range(1, 10).Select(lambda j: j + 1))
+    assert "function_call" not in ast.unparse(r)