refactor[next]: stricter typing in ffront

Author: havogt

pyproject.toml

@@ -206,6 +206,11 @@ warn_unused_ignores = false
 disallow_incomplete_defs = false
 module = 'gt4py.next.*'
 
+[[tool.mypy.overrides]]
+# TODO: temporarily to propagate it to all of next
+disallow_incomplete_defs = true
+module = 'gt4py.next.ffront.*'
+
 [[tool.mypy.overrides]]
 ignore_errors = true
 module = 'gt4py.next.ffront.decorator'

src/gt4py/next/ffront/decorator.py

@@ -24,7 +24,7 @@
 import typing
 import warnings
 from collections.abc import Callable
-from typing import Generator, Generic, TypeVar
+from typing import Generic, TypeVar
 
 from gt4py import eve
 from gt4py._core import definitions as core_defs
@@ -72,25 +72,6 @@
 DEFAULT_BACKEND: Callable = None
 
 
-def _field_constituents_shape_and_dims(
-    arg, arg_type: ts.FieldType | ts.ScalarType | ts.TupleType
-) -> Generator[tuple[tuple[int, ...], list[Dimension]]]:
-    if isinstance(arg_type, ts.TupleType):
-        for el, el_type in zip(arg, arg_type.types):
-            yield from _field_constituents_shape_and_dims(el, el_type)
-    elif isinstance(arg_type, ts.FieldType):
-        dims = type_info.extract_dims(arg_type)
-        if hasattr(arg, "shape"):
-            assert len(arg.shape) == len(dims)
-            yield (arg.shape, dims)
-        else:
-            yield (None, dims)
-    elif isinstance(arg_type, ts.ScalarType):
-        yield (None, [])
-    else:
-        raise ValueError("Expected 'FieldType' or 'TupleType' thereof.")
-
-
 # TODO(tehrengruber): Decide if and how programs can call other programs. As a
 #  result Program could become a GTCallable.
 @dataclasses.dataclass(frozen=True)
@@ -179,7 +160,7 @@ def with_backend(self, backend: ppi.ProgramExecutor) -> Program:
     def with_grid_type(self, grid_type: GridType) -> Program:
         return dataclasses.replace(self, grid_type=grid_type)
 
-    def with_bound_args(self, **kwargs) -> ProgramWithBoundArgs:
+    def with_bound_args(self, **kwargs: Any) -> ProgramWithBoundArgs:
         """
         Bind scalar, i.e. non field, program arguments.
 
@@ -229,7 +210,7 @@ def itir(self) -> itir.FencilDefinition:
             )
         ).program
 
-    def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs) -> None:
+    def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs: Any) -> None:
         if self.backend is None:
             warnings.warn(
                 UserWarning(
@@ -376,7 +357,9 @@ def program(
 
     def program_inner(definition: types.FunctionType) -> Program:
         return Program.from_function(
-            definition, DEFAULT_BACKEND if backend is eve.NOTHING else backend, grid_type
+            definition,
+            DEFAULT_BACKEND if backend is eve.NOTHING else backend,
+            grid_type,
         )
 
     return program_inner if definition is None else program_inner(definition)
@@ -634,9 +617,13 @@ def field_operator(definition=None, *, backend=eve.NOTHING, grid_type=None):
         ...     ...
     """
 
-    def field_operator_inner(definition: types.FunctionType) -> FieldOperator[foast.FieldOperator]:
+    def field_operator_inner(
+        definition: types.FunctionType,
+    ) -> FieldOperator[foast.FieldOperator]:
         return FieldOperator.from_function(
-            definition, DEFAULT_BACKEND if backend is eve.NOTHING else backend, grid_type
+            definition,
+            DEFAULT_BACKEND if backend is eve.NOTHING else backend,
+            grid_type,
         )
 
     return field_operator_inner if definition is None else field_operator_inner(definition)

src/gt4py/next/ffront/dialect_parser.py

@@ -84,7 +84,7 @@ def apply(
         return output_ast
 
     @classmethod
-    def apply_to_function(cls, function: Callable):
+    def apply_to_function(cls, function: Callable) -> DialectRootT:
         src = SourceDefinition.from_function(function)
         closure_vars = get_closure_vars_from_function(function)
         annotations = typing.get_type_hints(function)
@@ -96,7 +96,10 @@ def _preprocess_definition_ast(cls, definition_ast: ast.AST) -> ast.AST:
 
     @classmethod
     def _postprocess_dialect_ast(
-        cls, output_ast: DialectRootT, closure_vars: dict[str, Any], annotations: dict[str, Any]
+        cls,
+        output_ast: DialectRootT,
+        closure_vars: dict[str, Any],
+        annotations: dict[str, Any],
     ) -> DialectRootT:
         return output_ast
 

src/gt4py/next/ffront/foast_passes/closure_var_folding.py

@@ -35,12 +35,18 @@ class ClosureVarFolding(NodeTranslator, traits.VisitorWithSymbolTableTrait):
 
     @classmethod
     def apply(
-        cls, node: foast.FunctionDefinition | foast.FieldOperator, closure_vars: dict[str, Any]
+        cls,
+        node: foast.FunctionDefinition | foast.FieldOperator,
+        closure_vars: dict[str, Any],
     ) -> foast.FunctionDefinition:
         return cls(closure_vars=closure_vars).visit(node)
 
     def visit_Name(
-        self, node: foast.Name, current_closure_vars, symtable, **kwargs
+        self,
+        node: foast.Name,
+        current_closure_vars: dict[str, Any],
+        symtable: dict[str, foast.Symbol],
+        **kwargs: Any,
     ) -> foast.Name | foast.Constant:
         if node.id in symtable:
             definition = symtable[node.id]
@@ -50,7 +56,7 @@ def visit_Name(
                     return foast.Constant(value=value, location=node.location)
         return node
 
-    def visit_Attribute(self, node: foast.Attribute, **kwargs) -> foast.Constant:
+    def visit_Attribute(self, node: foast.Attribute, **kwargs: Any) -> foast.Constant:
         value = self.visit(node.value, **kwargs)
         if isinstance(value, foast.Constant):
             if hasattr(value.value, node.attr):
@@ -59,6 +65,6 @@ def visit_Attribute(self, node: foast.Attribute, **kwargs) -> foast.Constant:
         raise errors.DSLError(node.location, "Attribute access only applicable to constants.")
 
     def visit_FunctionDefinition(
-        self, node: foast.FunctionDefinition, **kwargs
+        self, node: foast.FunctionDefinition, **kwargs: Any
     ) -> foast.FunctionDefinition:
         return self.generic_visit(node, current_closure_vars=node.closure_vars, **kwargs)

src/gt4py/next/ffront/foast_passes/closure_var_type_deduction.py

@@ -39,7 +39,7 @@ def apply(
         return cls(closure_vars=closure_vars).visit(node)
 
     def visit_FunctionDefinition(
-        self, node: foast.FunctionDefinition, **kwargs
+        self, node: foast.FunctionDefinition, **kwargs: Any
     ) -> foast.FunctionDefinition:
         new_closure_vars: list[foast.Symbol] = []
         for sym in node.closure_vars:

src/gt4py/next/ffront/foast_passes/dead_closure_var_elimination.py

@@ -15,10 +15,10 @@
 from typing import Any
 
 import gt4py.next.ffront.field_operator_ast as foast
-from gt4py.eve import NodeTranslator, traits
+from gt4py import eve
 
 
-class DeadClosureVarElimination(NodeTranslator, traits.VisitorWithSymbolTableTrait):
+class DeadClosureVarElimination(eve.NodeTranslator, eve.traits.VisitorWithSymbolTableTrait):
     """Remove closure variable symbols that are not referenced in the AST."""
 
     _referenced_symbols: list[foast.Symbol]
@@ -27,7 +27,9 @@ class DeadClosureVarElimination(NodeTranslator, traits.VisitorWithSymbolTableTra
     def apply(cls, node: foast.FunctionDefinition) -> foast.FunctionDefinition:
         return cls().visit(node)
 
-    def visit_Name(self, node: foast.Name, symtable, **kwargs: Any) -> foast.Name:
+    def visit_Name(
+        self, node: foast.Name, symtable: dict[str, foast.Symbol], **kwargs: Any
+    ) -> foast.Name:
         if node.id in symtable:
             self._referenced_symbols.append(symtable[node.id])
         return node

src/gt4py/next/ffront/foast_passes/iterable_unpack.py

@@ -57,7 +57,7 @@ def _unique_tuple_symbol(self, node: foast.TupleTargetAssign) -> foast.Symbol[An
         self.unique_tuple_symbol_id += 1
         return sym
 
-    def visit_BlockStmt(self, node: foast.BlockStmt, **kwargs) -> foast.BlockStmt:
+    def visit_BlockStmt(self, node: foast.BlockStmt, **kwargs: Any) -> foast.BlockStmt:
         unrolled_stmts: list[foast.Assign | foast.BlockStmt | foast.Return] = []
 
         for stmt in node.stmts:
@@ -79,7 +79,10 @@ def visit_BlockStmt(self, node: foast.BlockStmt, **kwargs) -> foast.BlockStmt:
                         slice_indices = list(range(lower, upper))
                         tuple_slice = [
                             foast.Subscript(
-                                value=tuple_name, index=i, type=el_type, location=stmt.location
+                                value=tuple_name,
+                                index=i,
+                                type=el_type,
+                                location=stmt.location,
                             )
                             for i in slice_indices
                         ]
@@ -98,7 +101,10 @@ def visit_BlockStmt(self, node: foast.BlockStmt, **kwargs) -> foast.BlockStmt:
                         new_assign = foast.Assign(
                             target=subtarget,
                             value=foast.Subscript(
-                                value=tuple_name, index=index, type=el_type, location=stmt.location
+                                value=tuple_name,
+                                index=index,
+                                type=el_type,
+                                location=stmt.location,
                             ),
                             location=stmt.location,
                         )

src/gt4py/next/ffront/foast_passes/type_alias_replacement.py

@@ -37,7 +37,9 @@ class TypeAliasReplacement(NodeTranslator, traits.VisitorWithSymbolTableTrait):
 
     @classmethod
     def apply(
-        cls, node: foast.FunctionDefinition | foast.FieldOperator, closure_vars: dict[str, Any]
+        cls,
+        node: foast.FunctionDefinition | foast.FieldOperator,
+        closure_vars: dict[str, Any],
     ) -> tuple[foast.FunctionDefinition, dict[str, Any]]:
         foast_node = cls(closure_vars=closure_vars).visit(node)
         new_closure_vars = closure_vars.copy()
@@ -53,10 +55,12 @@ def is_type_alias(self, node_id: SymbolName | SymbolRef) -> bool:
             and node_id not in TYPE_BUILTIN_NAMES
         )
 
-    def visit_Name(self, node: foast.Name, **kwargs) -> foast.Name:
+    def visit_Name(self, node: foast.Name, **kwargs: Any) -> foast.Name:
         if self.is_type_alias(node.id):
             return foast.Name(
-                id=self.closure_vars[node.id].__name__, location=node.location, type=node.type
+                id=self.closure_vars[node.id].__name__,
+                location=node.location,
+                type=node.type,
             )
         return node
 
@@ -79,7 +83,8 @@ def _update_closure_var_symbols(
                                 kw_only_args={},
                                 pos_only_args=[ts.DeferredType(constraint=ts.ScalarType)],
                                 returns=cast(
-                                    ts.DataType, from_type_hint(self.closure_vars[var.id])
+                                    ts.DataType,
+                                    from_type_hint(self.closure_vars[var.id]),
                                 ),
                             ),
                             namespace=dialect_ast_enums.Namespace.CLOSURE,
@@ -94,7 +99,7 @@ def _update_closure_var_symbols(
         return new_closure_vars
 
     def visit_FunctionDefinition(
-        self, node: foast.FunctionDefinition, **kwargs
+        self, node: foast.FunctionDefinition, **kwargs: Any
     ) -> foast.FunctionDefinition:
         return foast.FunctionDefinition(
             id=node.id,