Repository: iris-hep/func_adl

Issue:
Cannot use imported module names in lambdas
These lines don't check that a variable actually represents a constant before replacing a `Name` node with a `Constant` node:

https://github.com/iris-hep/func_adl/blob/89713708d8d1dc7234a38d78a6961170f4d26dca/func_adl/util_ast.py#L12-L15

For Python versions before 3.8, some type checking is done:

https://github.com/iris-hep/func_adl/blob/89713708d8d1dc7234a38d78a6961170f4d26dca/func_adl/util_ast.py#L17-L29

but in either case, this behavior causes issues here:

https://github.com/iris-hep/func_adl/blob/89713708d8d1dc7234a38d78a6961170f4d26dca/func_adl/util_ast.py#L335-L336

when `v` does not represent a literal. This causes problems for the Uproot backend, where I allow references to modules like `np` and `ak`. For example, I might use something like `ds.Select(lambda event: np.cos(event.MET_phi))`. That works fine as long as you haven't imported the relevant module in the source where the lambda is (since there is no potential capture), but it breaks if that module has been imported. If the surrounding code needs that module, the only solution is to go back to passing the lambda as a string.

I would argue that the correct behavior here is for the unmodified `Name` node to be passed through if the captured variable is not a constant (maybe with a warning or even just a debug log message). That's consistent with what happens if the name hasn't been defined outside of the lambda. This leaves it up to the particular code generator how to handle any apparently undefined `Name`, which basically allows each backend to set its own reserved keywords.

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/iris-hep/func_adl/pull/131?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`7295efc`)](https://app.codecov.io/gh/iris-hep/func_adl/commit/7295efc5dfee9c731fc21665abc54eb1c95e380b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) 96.79% compared to head [(`2c4e2c7`)](https://app.codecov.io/gh/iris-hep/func_adl/pull/131?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) 96.80%.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master     #131   +/-   ##
=======================================
  Coverage   96.79%   96.80%           
=======================================
  Files          15       15           
  Lines        1281     1283    +2     
=======================================
+ Hits         1240     1242    +2     
  Misses         41       41           
```

| [Flag](https://app.codecov.io/gh/iris-hep/func_adl/pull/131/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/iris-hep/func_adl/pull/131/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) | `96.80% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Files](https://app.codecov.io/gh/iris-hep/func_adl/pull/131?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) | Coverage Δ | |
|---|---|---|
| [func\_adl/util\_ast.py](https://app.codecov.io/gh/iris-hep/func_adl/pull/131?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep#diff-ZnVuY19hZGwvdXRpbF9hc3QucHk=) | `92.24% <100.00%> (+0.06%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iris-hep/func_adl/pull/131?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep).

masonproffitt: The problem of `as_literal` being called on non-literals is still there. For example, if you say `np = (i for i in [1])` somewhere before the lambda, this is still going to give you `ast.Constant(value=<generator object <genexpr> at [...]>)`, which is not a valid Python AST node. (Maybe a contrived example, but this AST would make debugging such a case much more difficult than it should be.)
masonproffitt: As I mentioned in the issue (#127), this type checking is already done by `func-adl` for Python versions < 3.8, so the suggestion is just to do this for Python >= 3.8 as well.
gordonwatts: Ok, @masonproffitt - I've added a whitelist for types we can capture - the error should be more informative. For example (from the test):

```python
ValueError("Do not know how to capture data type \'bogus\' for variable \'my_var\' - only str, int, float, bool, complex, str, bytes are supported."
```

which happens from this test:

```python
def test_sensible_error_with_bad_variable_capture():
    class bogus:
        def __init__(self):
            self.my_var = 10

    my_var = bogus()

    with pytest.raises(ValueError) as e:
        parse_as_ast(lambda x: x > my_var)

    assert "my_var" in str(e)
```

Does that get to what you are after?

Patch:
diff --git a/func_adl/util_ast.py b/func_adl/util_ast.py
index 5c5f8a4..ab43f47 100644
--- a/func_adl/util_ast.py
+++ b/func_adl/util_ast.py
@@ -5,6 +5,7 @@ import inspect
 import sys
 import tokenize
 from collections import defaultdict
+from types import ModuleType
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union, cast
 
 # Some functions to enable backwards compatibility.
@@ -333,7 +334,18 @@ class _rewrite_captured_vars(ast.NodeTransformer):
         if node.id in self._lookup_dict:
             v = self._lookup_dict[node.id]
             if not callable(v):
-                return as_literal(self._lookup_dict[node.id])
+                # Modules should be sent on down to be dealt with by the
+                # backend.
+                if not isinstance(v, ModuleType):
+                    legal_capture_types = [str, int, float, bool, complex, str, bytes]
+                    if type(v) not in legal_capture_types:
+                        raise ValueError(
+                            f"Do not know how to capture data type '{type(v).__name__}' for "
+                            f"variable '{node.id}' - only "
+                            f"{', '.join([c.__name__ for c in legal_capture_types])} are "
+                            "supported."
+                        )
+                    return as_literal(v)
         return node
 
     def visit_Lambda(self, node: ast.Lambda) -> Any:
diff --git a/setup.py b/setup.py
index 67668bc..208056e 100644
--- a/setup.py
+++ b/setup.py
@@ -21,6 +21,7 @@ extras_require = {
         "astunparse",
         "black",
         "isort",
+        "numpy",
     ]
 }
 extras_require["complete"] = sorted(set(sum(extras_require.values(), [])))


Test patch:
diff --git a/tests/test_object_stream.py b/tests/test_object_stream.py
index e4c3b85..492018d 100644
--- a/tests/test_object_stream.py
+++ b/tests/test_object_stream.py
@@ -35,9 +35,7 @@ class dd_jet:
 T = TypeVar("T")
 
 
-def add_md_for_type(
-    s: ObjectStream[T], a: ast.Call
-) -> Tuple[ObjectStream[T], ast.Call]:
+def add_md_for_type(s: ObjectStream[T], a: ast.Call) -> Tuple[ObjectStream[T], ast.Call]:
     return s.MetaData({"hi": "there"}), a
 
 
@@ -323,9 +321,7 @@ def test_query_metadata_not_empty():
 def test_nested_query_rendered_correctly():
     r = (
         my_event()
-        .Where(
-            "lambda e: e.jets.Select(lambda j: j.pT()).Where(lambda j: j > 10).Count() > 0"
-        )
+        .Where("lambda e: e.jets.Select(lambda j: j.pT()).Where(lambda j: j > 10).Count() > 0")
         .SelectMany("lambda e: e.jets()")
         .AsROOTTTree("junk.root", "analysis", "jetPT")
         .value()
@@ -346,6 +342,25 @@ def test_query_with_comprehensions():
     assert "ListComp" not in ast.dump(r)
 
 
+def test_non_imported_function_call():
+    r = (
+        my_event()
+        .Select(lambda event: np.cos(event.MET_phi))  # type: ignore # noqa
+        .Where(lambda p: p > 0.0)
+        .value()
+    )  # NOQA
+    assert isinstance(r, ast.AST)
+    assert "Attribute(value=Name(id='np', ctx=Load()), attr='cos', ctx=Load())" in ast.dump(r)
+
+
+def test_imported_function_call():
+    import numpy as np
+
+    r = my_event().Select(lambda event: np.cos(event.MET_phi)).Where(lambda p: p > 0.0).value()
+    assert isinstance(r, ast.AST)
+    assert "Attribute(value=Name(id='np', ctx=Load()), attr='cos', ctx=Load())" in ast.dump(r)
+
+
 def test_bad_where():
     with pytest.raises(ValueError):
         my_event().Where("lambda e: 10").SelectMany("lambda e: e.jets()").AsROOTTTree(
diff --git a/tests/test_util_ast.py b/tests/test_util_ast.py
index e4a8ee1..cf54715 100644
--- a/tests/test_util_ast.py
+++ b/tests/test_util_ast.py
@@ -309,6 +309,19 @@ def test_parse_lambda_capture_nested_local():
     assert ast.dump(r) == ast.dump(r_true)
 
 
+def test_sensible_error_with_bad_variable_capture():
+    class bogus:
+        def __init__(self):
+            self.my_var = 10
+
+    my_var = bogus()
+
+    with pytest.raises(ValueError) as e:
+        parse_as_ast(lambda x: x > my_var)
+
+    assert "my_var" in str(e)
+
+
 def test_parse_simple_func():
     "A oneline function defined at local scope"