Repository: iris-hep/func_adl

Issue:
Zip lookup for function types is causing a type tracing failure
This is from @ponyisi. 

You get this error:

```text
  File "/home/onyisi/servicex/analysis-grand-challenge/analyses/cms-open-data-ttbar/venv/lib64/python3.9/site-packages/func_adl/type_based_replacement.py", line 971, in visit_Attribute
    raise ValueError(f"Key {key} not found in dict expression!!")
ValueError: Key Zip not found in dict expression!!
```

When running this AGC code:

```python
    cuts = source.Where(lambda e: {"pt": e.Electron_pt,
                               "eta": e.Electron_eta,
                               "cutBased": e.Electron_cutBased,
                               "sip3d": e.Electron_sip3d,}.Zip()\
                        .Where(lambda electron: (electron.pt > 30
                                                 and abs(electron.eta) < 2.1
                                                 and electron.cutBased == 4
                                                 and electron.sip3d < 4)).Count()
                        + {"pt": e.Muon_pt,
                           "eta": e.Muon_eta,
                           "tightId": e.Muon_tightId,
                           "sip3d": e.Muon_sip3d,
                           "pfRelIso04_all": e.Muon_pfRelIso04_all}.Zip()\
                        .Where(lambda muon: (muon.pt > 30
                                             and abs(muon.eta) < 2.1
                                             and muon.tightId
                                             and muon.pfRelIso04_all < 0.15)).Count()== 1)\
                        .Where(lambda f: {"pt": f.Jet_pt,
                                          "eta": f.Jet_eta,
                                          "jetId": f.Jet_jetId}.Zip()\
                               .Where(lambda jet: (jet.pt > 25
                                                   and abs(jet.eta) < 2.4
                                                   and jet.jetId == 6)).Count() >= 4)\
                        .Where(lambda g: {"pt": g.Jet_pt,
                                          "eta": g.Jet_eta,
                                          "btagCSVV2": g.Jet_btagCSVV2,
                                          "jetId": g.Jet_jetId}.Zip()\
                        .Where(lambda jet: (jet.btagCSVV2 > 0.5
                                            and jet.pt > 25
                                            and abs(jet.eta) < 2.4)
                                            and jet.jetId == 6).Count() >= 1)
```

Likely the problem is the function lookup should trigger an `Any` type return, not a total failure like this. And, further, it should work if the system has some type information (to catch mispellings).

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/iris-hep/func_adl/pull/149?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) Report
Attention: Patch coverage is `75.00000%` with `3 lines` in your changes missing coverage. Please review.
> Project coverage is 96.93%. Comparing base [(`95b14c9`)](https://app.codecov.io/gh/iris-hep/func_adl/commit/95b14c9b3986201236d4f812dd4499185e99721e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) to head [(`77e4888`)](https://app.codecov.io/gh/iris-hep/func_adl/commit/77e4888535c469ab45e9f0a447279a11f2a2fdeb?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep).

| [Files with missing lines](https://app.codecov.io/gh/iris-hep/func_adl/pull/149?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) | Patch % | Lines |
|---|---|---|
| [func\_adl/ast/function\_simplifier.py](https://app.codecov.io/gh/iris-hep/func_adl/pull/149?src=pr&el=tree&filepath=func_adl%2Fast%2Ffunction_simplifier.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep#diff-ZnVuY19hZGwvYXN0L2Z1bmN0aW9uX3NpbXBsaWZpZXIucHk=) | 40.00% | [3 Missing :warning: ](https://app.codecov.io/gh/iris-hep/func_adl/pull/149?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) |

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


```diff
@@           Coverage Diff           @@
##           master     #149   +/-   ##
=======================================
  Coverage   96.93%   96.93%           
=======================================
  Files          15       15           
  Lines        1336     1338    +2     
=======================================
+ Hits         1295     1297    +2     
  Misses         41       41           
```

| [Flag](https://app.codecov.io/gh/iris-hep/func_adl/pull/149/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/149/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=iris-hep) | `96.93% <75.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.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/iris-hep/func_adl/pull/149?dropdown=coverage&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).


Patch:
diff --git a/func_adl/ast/function_simplifier.py b/func_adl/ast/function_simplifier.py
index a476799..c6f63c6 100644
--- a/func_adl/ast/function_simplifier.py
+++ b/func_adl/ast/function_simplifier.py
@@ -244,8 +244,8 @@ class simplify_chained_calls(FuncADLNodeTransformer):
             "SelectMany", [cast(ast.AST, captured_body), cast(ast.AST, func_g)]
         )
         new_select_lambda = lambda_build(captured_arg, new_select)
-        new_selectmany = function_call("SelectMany", [seq, cast(ast.AST, new_select_lambda)])
-        return new_selectmany
+        new_select_many = function_call("SelectMany", [seq, cast(ast.AST, new_select_lambda)])
+        return new_select_many
 
     def call_SelectMany(self, node: ast.Call, args: List[ast.AST]):
         r"""
@@ -442,7 +442,7 @@ class simplify_chained_calls(FuncADLNodeTransformer):
         # Get the value out - this is due to supporting python 3.7-3.9
         n = _get_value_from_index(s)
         if n is None:
-            return ast.Subscript(v, s, ast.Load())
+            return ast.Subscript(v, s, ast.Load())  # type: ignore
         assert isinstance(n, int), "Programming error: index is not an integer in tuple subscript"
         if n >= len(v.elts):
             raise FuncADLIndexError(
@@ -460,7 +460,7 @@ class simplify_chained_calls(FuncADLNodeTransformer):
         """
         n = _get_value_from_index(s)
         if n is None:
-            return ast.Subscript(v, s, ast.Load())
+            return ast.Subscript(v, s, ast.Load())  # type: ignore
         if n >= len(v.elts):
             raise FuncADLIndexError(
                 f"Attempt to access the {n}th element of a tuple"
@@ -484,7 +484,7 @@ class simplify_chained_calls(FuncADLNodeTransformer):
             if _get_value_from_index(value) == s:
                 return v.values[index]
 
-        return ast.Subscript(v, s, ast.Load())
+        return ast.Subscript(v, s, ast.Load())  # type: ignore
 
     def visit_Subscript_Of_First(self, first: ast.AST, s):
         """
diff --git a/func_adl/type_based_replacement.py b/func_adl/type_based_replacement.py
index 1ccc51d..cfc7563 100644
--- a/func_adl/type_based_replacement.py
+++ b/func_adl/type_based_replacement.py
@@ -102,7 +102,7 @@ def func_adl_callable(
         Callable[[ObjectStream[W], ast.Call], Tuple[ObjectStream[W], ast.AST]]
     ] = None
 ):
-    """Dectorator that will declare a function that can be used inline in
+    """Decorator that will declare a function that can be used inline in
     a `func_adl` expression. The body of the function, what the backend
     translates it to, must be given by another route (e.g. via `MetaData`
     and the `processor` argument).
@@ -388,7 +388,7 @@ def _fill_in_default_arguments(func: Callable, call: ast.Call) -> Tuple[ast.Call
         )
         return_type = Any
 
-    return call, return_type
+    return call, return_type  # type: ignore
 
 
 def fixup_ast_from_modifications(transformed_ast: ast.AST, original_ast: ast.Call) -> ast.Call:
@@ -582,7 +582,7 @@ def remap_by_types(
             # If this is a known collection class, we can use call-backs to follow it.
             if get_origin(call_site_info.obj_type) in _g_collection_classes:
                 rtn_value = self.process_method_call_on_stream_obj(
-                    _g_collection_classes[get_origin(call_site_info.obj_type)],
+                    _g_collection_classes[get_origin(call_site_info.obj_type)],  # type: ignore
                     m_name,
                     r_node,
                     get_args(call_site_info.obj_type)[0],
@@ -631,7 +631,7 @@ def remap_by_types(
 
             Args:
                 node (ast.Call): The ast node
-                obj_type (type): The object type this method call is occuring against
+                obj_type (type): The object type this method call is occurring against
 
             Returns:
                 ast.AST: An updated ast that is the new method call (with default args, etc.)
@@ -643,7 +643,7 @@ def remap_by_types(
             base_obj_list_all = [obj_type]
             if is_iterable(obj_type):
                 item_type = unwrap_iterable(obj_type)
-                base_obj_list_all += [c[item_type] for c in _g_collection_classes]
+                base_obj_list_all += [c[item_type] for c in _g_collection_classes]  # type: ignore
 
             assert isinstance(r_node.func, ast.Attribute)
             m_name = r_node.func.attr
@@ -758,7 +758,7 @@ def remap_by_types(
                     f"function {func_info.function.__name__} ({str(e)})"
                 ) from e
 
-        def process_paramaterized_method_call(
+        def process_parameterized_method_call(
             self,
             node: ast.Call,
             obj_type: Type,
@@ -815,7 +815,7 @@ def remap_by_types(
                 if isinstance(t_node.func.value, ast.Attribute):
                     found_type = self.lookup_type(t_node.func.value.value)
                     if found_type is not None:
-                        t_node = self.process_paramaterized_method_call(
+                        t_node = self.process_parameterized_method_call(
                             t_node,
                             found_type,
                             t_node.func.value.attr,
@@ -941,17 +941,17 @@ def remap_by_types(
             return node
 
         def visit_Num(self, node: ast.Num) -> Any:  # pragma: no cover
-            "3.7 compatability"
+            "3.7 compatibility"
             self._found_types[node] = type(node.n)
             return node
 
         def visit_Str(self, node: ast.Str) -> Any:  # pragma: no cover
-            "3.7 compatability"
+            "3.7 compatibility"
             self._found_types[node] = str
             return node
 
         def visit_NameConstant(self, node: ast.NameConstant) -> Any:  # pragma: no cover
-            "3.7 compatability"
+            "3.7 compatibility"
             if node.value is None:
                 raise ValueError("Do not know how to work with pythons None")
             self._found_types[node] = bool
@@ -968,6 +968,8 @@ def remap_by_types(
                     e for e, k in enumerate(t_node.value.keys) if k.value == key  # type: ignore
                 ]
                 if len(key_index) == 0:
+                    if t_node.attr.lower() == "zip":
+                        return t_node
                     raise ValueError(f"Key {key} not found in dict expression!!")
                 value = t_node.value.values[key_index[0]]
                 self._found_types[node] = self.lookup_type(value)
@@ -999,7 +1001,7 @@ def remap_from_lambda(
     orig_type = o_stream.item_type
     var_name = l_func.args.args[0].arg
     stream, new_body, return_type = remap_by_types(o_stream, var_name, orig_type, l_func.body)
-    return stream, ast.Lambda(l_func.args, new_body), return_type
+    return stream, ast.Lambda(l_func.args, new_body), return_type  # type: ignore
 
 
 def reset_global_functions():


Test patch:
diff --git a/tests/test_type_based_replacement.py b/tests/test_type_based_replacement.py
index ca39992..a2641f5 100644
--- a/tests/test_type_based_replacement.py
+++ b/tests/test_type_based_replacement.py
@@ -546,6 +546,23 @@ def test_dictionary_bad_key():
     assert "jetsss" in str(e)
 
 
+def test_dictionary_Zip_key():
+    "Check that type follow from a dictionary through a Zip works"
+
+    s = ast_lambda(
+        """({
+        'jet_pt':  e.Jets().Select(lambda j: j.pt()),
+        'jet_eta': e.Jets().Select(lambda j: j.eta())}
+        .Zip()
+        .Select(lambda j: j.pt()))"""
+    )
+    objs = ObjectStream[Event](ast.Name(id="e", ctx=ast.Load()))
+
+    new_objs, new_s, expr_type = remap_by_types(objs, "e", Event, s)
+
+    assert expr_type == Any
+
+
 def test_dictionary_through_Select():
     """Make sure the Select statement carries the typing all the way through"""