Type propagation through dict and tuples (#140)

* Initial go at the tests for this

* Deal with dictionaries and tuples for type following

* Fix up flake8 and type errors

* Update github action references

* Deal with python 3.8

Author: gordonwatts

.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

func_adl/type_based_replacement.py

@@ -496,7 +496,16 @@ def stream(self) -> ObjectStream[S]:
 
         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 visit_IfExp(self, node: ast.IfExp) -> Any:
         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 visit_NameConstant(self, node: ast.NameConstant) -> Any:  # pragma: no cover
             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)
 

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 @@ class object where the method is defined, and the method object. If there is no
     # 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 @@ class object where the method is defined, and the method object. If there is no
             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

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)