Repository: iris-hep/func_adl

Issue:
Typing is not followed through dictionaries
It would seem that type checking is not followed through dictionaries:

```python
query = (ds
         .Select(lambda e: {
             'evt': e.EventInfo("EventInfo"),
             'jet': e.Jets()
             })
         .Select(lambda ei: {
             'event_number': ei.evt.eventNumber(),
             'run_number': ei.evt.runNumber(),
             'jet_pt': ei.jet.Select(lambda j: j.pt()/1000),
             'jet_energy_sampling': ei.jet.Select(lambda j: j.getAttribute[cpp_vfloat]("EnergyPerSampling")),
         })
)
```

In this case, the `getAttribute` is not processed as a parameterized call because the type of `ei.jet` is not known! The place this seems to happen is here:

```python
        def process_method_call(self, node: ast.Call, obj_type: type) -> ast.AST:
            """Manage a call against a typed object.

            1. Fill in default arguments based on the signature from the object type.
            1. Insert metadata into the stream if there are call backs based on the object type.
            1. Type follow the call site to get the return type, and to make sure anything inside
               lambda's etc., get correctly type followed as well.

            Args:
                node (ast.Call): The ast node
                obj_type (type): The object type this method call is occuring against

            Returns:
                ast.AST: An updated ast that is the new method call (with default args, etc.)
            """
            # Make reference copies that we'll populate as we go
            r_node = node

            # Find all objects that this method exists on.
            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]

            assert isinstance(r_node.func, ast.Attribute)

```

The type comes in as `Any` for this call.

This was discovered as a part of the [idap challenge](https://github.com/iris-hep/idap-200gbps-atlas/issues/47)

Discussion:

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 98869bd..180aa69 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -12,9 +12,9 @@ jobs:
   flake8:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python 3.8
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: 3.8
       - name: Install dependencies
@@ -41,9 +41,9 @@ jobs:
         python-version: [3.8, 3.9, "3.10", 3.11, 3.12]
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
diff --git a/func_adl/type_based_replacement.py b/func_adl/type_based_replacement.py
index 6727ea9..45b524a 100644
--- a/func_adl/type_based_replacement.py
+++ b/func_adl/type_based_replacement.py
@@ -496,7 +496,16 @@ def remap_by_types(
 
         def lookup_type(self, name: Union[str, object]) -> Type:
             "Return the type for a node, Any if we do not know about it"
-            return self._found_types.get(name, Any)  # type: ignore
+            t = self._found_types.get(name, None)
+            if t is not None:
+                return t  # type: ignore
+            if not isinstance(name, ast.AST):
+                return Any  # type: ignore
+
+            # It could be we can determine the type from this ast.
+            return Any  # type: ignore
+
+            return t  # type: ignore
 
         def process_method_call_on_stream_obj(
             self,
@@ -883,9 +892,25 @@ def remap_by_types(
         def visit_Subscript(self, node: ast.Subscript) -> Any:
             t_node = self.generic_visit(node)
             assert isinstance(t_node, ast.Subscript)
-            inner_type = unwrap_iterable(self.lookup_type(t_node.value))
-            self._found_types[node] = inner_type
-            self._found_types[t_node] = inner_type
+            if isinstance(t_node.value, ast.Tuple):
+                slice = t_node.slice
+                # This if statement can be removed when we no longer support 3.8.
+                if isinstance(slice, ast.Index):
+                    slice = slice.value  # type: ignore
+                if not isinstance(slice, ast.Constant):
+                    raise ValueError(
+                        f"Slices must be indexable constants only - {ast.dump(slice)} is not "
+                        "valid."
+                    )
+                index = slice.value
+                if len(t_node.value.elts) <= index:
+                    raise ValueError(f"Index {index} out of range for {ast.dump(node.value)}")
+                self._found_types[node] = self.lookup_type(t_node.value.elts[index])
+                self._found_types[t_node] = self.lookup_type(t_node.value.elts[index])
+            else:
+                inner_type = unwrap_iterable(self.lookup_type(t_node.value))
+                self._found_types[node] = inner_type
+                self._found_types[t_node] = inner_type
             return t_node
 
         def visit_Name(self, node: ast.Name) -> ast.Name:
@@ -919,6 +944,22 @@ def remap_by_types(
             self._found_types[node] = bool
             return node
 
+        def visit_Attribute(self, node: ast.Attribute) -> Any:
+            t_node = self.generic_visit(node)
+            assert isinstance(t_node, ast.Attribute)
+            # If this is a dict reference, then figure out what the
+            # type is for that value of the dict.
+            if isinstance(t_node.value, ast.Dict):
+                key = t_node.attr
+                key_index = [
+                    e for e, k in enumerate(t_node.value.keys) if k.value == key  # type: ignore
+                ]
+                if len(key_index) == 0:
+                    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)
+            return t_node
+
     tt = type_transformer(o_stream)
     r_a = tt.visit(a)
 
diff --git a/func_adl/util_types.py b/func_adl/util_types.py
index 42eae29..a3fedc5 100644
--- a/func_adl/util_types.py
+++ b/func_adl/util_types.py
@@ -46,7 +46,7 @@ def get_inherited(t: Type) -> Type:
     elif hasattr(t, "__origin__") and hasattr(t.__origin__, "__orig_bases__"):
         base_classes = t.__origin__.__orig_bases__
     else:
-        return Any
+        return Any  # type: ignore
 
     r = base_classes[0]  # type: ignore
 
@@ -77,7 +77,7 @@ def unwrap_iterable(t: Type) -> Type:
         t = get_inherited(t)
 
     if t == Any:
-        return Any
+        return Any  # type: ignore
 
     a = get_args(t)
     assert len(a) == 1, f"Coding error - expected iterable type with a parameter, got {t}"
@@ -168,7 +168,7 @@ def resolve_type_vars(
         s = build_type_dict_from_type(context_type, at_class)
     except TypeError:
         s = {}
-    return _resolve_type(parameterized_type, s)
+    return _resolve_type(parameterized_type, s)  # type: ignore
 
 
 def get_class_name(t: Type) -> str:
@@ -206,7 +206,7 @@ def get_method_and_class(class_object: Type, method_name: str) -> Optional[Tuple
     # Check for templated classes
     # TODO: Use inspect.getmro
     if not hasattr(class_object, "__mro__"):
-        class_object = get_origin(class_object)
+        class_object = get_origin(class_object)  # type: ignore
 
     # Walk the resolution hierarchy to find the method
     found_obj = None
@@ -223,6 +223,6 @@ def get_method_and_class(class_object: Type, method_name: str) -> Optional[Tuple
             if found_method == m:
                 found_obj = c
             else:
-                return (found_obj, found_method)
+                return (found_obj, found_method)  # type: ignore
 
-    return (found_obj, found_method)
+    return (found_obj, found_method)  # type: ignore


Test patch:
diff --git a/tests/test_type_based_replacement.py b/tests/test_type_based_replacement.py
index 3b3f699..94ad0f5 100644
--- a/tests/test_type_based_replacement.py
+++ b/tests/test_type_based_replacement.py
@@ -503,6 +503,64 @@ def test_collection_Select(caplog):
     assert len(caplog.text) == 0
 
 
+def test_dictionary():
+    "Check that we can type-follow through dictionaries"
+
+    s = ast_lambda("{'jets': e.Jets()}.jets.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 == Iterable[float]
+
+
+def test_dictionary_bad_key():
+    "Check that we can type-follow through dictionaries"
+
+    with pytest.raises(ValueError) as e:
+        s = ast_lambda("{'jets': e.Jets()}.jetsss.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 "jetsss" in str(e)
+
+
+def test_indexed_tuple():
+    "Check that we can type-follow through dictionaries"
+
+    s = ast_lambda("(e.Jets(),)[0].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 == Iterable[float]
+
+
+def test_indexed_tuple_out_of_bounds():
+    "Check that we can type-follow through dictionaries"
+
+    with pytest.raises(ValueError) as e:
+        s = ast_lambda("(e.Jets(),)[3].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 "3" in str(e)
+
+
+def test_indexed_tuple_bad_slice():
+    "Check that we can type-follow through dictionaries"
+
+    with pytest.raises(ValueError) as e:
+        s = ast_lambda("(e.Jets(),)[0:1].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 "is not valid" in str(e)
+
+
 def test_collection_Select_meta(caplog):
     "A simple collection"
     caplog.set_level(logging.WARNING)