Mypy fixes ones typing has been activated

Author: jorgensd

ufl/algorithms/apply_restrictions.py

@@ -19,6 +19,7 @@
 from ufl.corealg.map_dag import map_expr_dag
 from ufl.corealg.multifunction import MultiFunction
 from ufl.domain import Mesh, extract_unique_domain
+from ufl.integral import Integral
 from ufl.sobolevspace import H1
 
 default_restriction_map = {
@@ -291,7 +292,9 @@ def facet_normal(self, o):
             return self._require_restriction(o)
 
 
-def apply_restrictions(expression: Expr, default_restrictions: dict | None = None) -> Expr:
+def apply_restrictions(
+    expression: Expr | Integral, default_restrictions: dict | None = None
+) -> Expr:
     """Propagate restriction nodes to wrap differential terminals directly.
 
     Args:

ufl/algorithms/domain_analysis.py

@@ -10,7 +10,10 @@
 
 import numbers
 from collections import defaultdict
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from ufl.classes import Coefficient
 
 import ufl
 from ufl.algorithms.coordinate_derivative_helpers import (
@@ -82,8 +85,8 @@ def __init__(
 
         # This is populated in preprocess using data not available at
         # this stage:
-        self.integral_coefficients = None
-        self.enabled_coefficients = None
+        self.integral_coefficients: set[Coefficient] | None = None
+        self.enabled_coefficients: list[bool] | None = None
         non_primal_domains = tuple(domain_integral_type_map.keys())[1:]
         if sort_domains(non_primal_domains) != non_primal_domains:
             raise ValueError("domain_integral_type_map must have been sorted by domains")
@@ -336,14 +339,14 @@ def keyfunc(item):
         )
 
     integral_datas = []
-    for (d, itype, sid, extra_d_itype_tuple), integrals in sorted(itgs.items(), key=keyfunc):
+    for (d, itype, sid, extra_d_itype_tuple), _itgs in sorted(itgs.items(), key=keyfunc):
         d_itype_tuple = ((d, itype),) + extra_d_itype_tuple
         integral_datas.append(
             IntegralData(
                 d,
                 itype,
                 sid,
-                integrals,
+                _itgs,
                 {},
                 dict(d_itype_tuple),
             )

ufl/algorithms/formdata.py

@@ -20,9 +20,8 @@
 from ufl.algorithms.domain_analysis import IntegralData, reconstruct_form_from_integral_data
 from ufl.algorithms.replace import replace
 from ufl.classes import Argument, Coefficient, FunctionSpace, GeometricFacetQuantity
-from ufl.coefficient import BaseCoefficient
 from ufl.corealg.traversal import traverse_unique_terminals
-from ufl.domain import MeshSequence, extract_domains, extract_unique_domain
+from ufl.domain import AbstractDomain, MeshSequence, extract_domains, extract_unique_domain
 from ufl.form import Form, Zero
 from ufl.utils.formatting import estr, lstr, tstr
 from ufl.utils.sequences import max_degree
@@ -99,7 +98,7 @@ def _compute_element_mapping(form: Form):
 
 def _compute_max_subdomain_ids(integral_data: list[IntegralData]) -> dict[str, int]:
     """Compute the maximum subdomain ids."""
-    max_subdomain_ids = {}
+    max_subdomain_ids: dict[str, int] = {}
     for itg_data in integral_data:
         it = itg_data.integral_type
         for integral in itg_data.integrals:
@@ -141,8 +140,8 @@ def _check_facet_geometry(integral_data):
 
 
 def _build_coefficient_replace_map(
-    coefficients: list[BaseCoefficient], element_mapping=None
-) -> tuple[list[BaseCoefficient], dict[BaseCoefficient, BaseCoefficient]]:
+    coefficients: list[Coefficient], element_mapping=None
+) -> tuple[list[Coefficient], dict[Coefficient, Coefficient]]:
     """Create new Coefficient objects with count starting at 0.
 
     Returns:
@@ -175,10 +174,10 @@ class FormData:
 
     _original_form: Form
     _integral_data: list[IntegralData]
-    _reduced_coefficients: list[BaseCoefficient]
-    _original_coefficient_positions: list[tuple[int, BaseCoefficient]]
-    _function_replace_map: dict[BaseCoefficient, BaseCoefficient]
-    _coefficient_elements: list[Any]
+    _reduced_coefficients: list[Coefficient]
+    _original_coefficient_positions: list[int]
+    _function_replace_map: dict[Coefficient, Coefficient]
+    _coefficient_elements: tuple[Any, ...]
     _coefficient_split: dict[Coefficient, list[Coefficient]]
 
     def __init__(
@@ -188,7 +187,7 @@ def __init__(
         do_apply_default_restrictions: bool = True,
         do_apply_restrictions: bool = True,
         do_replace_functions: bool = False,
-        coefficients_to_split: tuple[BaseCoefficient, ...] | None = None,
+        coefficients_to_split: tuple[Coefficient, ...] | None = None,
         complex_mode: bool = False,
     ):
         """Create form-data for a form that has been processed.
@@ -223,8 +222,9 @@ def __init__(
         # Figure out which coefficients from the original form are
         # actually used in any integral (Differentiation may reduce the
         # set of coefficients w.r.t. the original form)
-        reduced_coefficients_set = set()
+        reduced_coefficients_set: set[Coefficient] = set()
         for itg_data in self.integral_data:
+            assert itg_data.integral_coefficients is not None
             reduced_coefficients_set.update(itg_data.integral_coefficients)
         self._reduced_coefficients = sorted(reduced_coefficients_set, key=lambda c: c.count())
         self._original_coefficient_positions = [
@@ -236,6 +236,7 @@ def __init__(
         # Store back into integral data which form coefficients are used
         # by each integral
         for itg_data in self.integral_data:
+            assert itg_data.integral_coefficients is not None
             itg_data.enabled_coefficients = [
                 bool(coeff in itg_data.integral_coefficients) for coeff in self.reduced_coefficients
             ]
@@ -285,7 +286,7 @@ def __init__(
                     elem = c.ufl_element()
                     coefficient_split[c] = [
                         Coefficient(FunctionSpace(m, e))
-                        for m, e in zip(mesh.iterable_like(elem), elem.sub_elements)
+                        for m, e in zip(mesh.iterable_like(elem), elem.sub_elements)  # type: ignore
                     ]
             self._coefficient_split = coefficient_split
             coeff_splitter = CoefficientSplitter(self.coefficient_split)
@@ -309,6 +310,7 @@ def __init__(
                     for _, integral_type in itg_data.domain_integral_type_map.items()
                 ):
                     continue
+                default_restrictions: dict[AbstractDomain, str | None] | None
                 if do_apply_default_restrictions:
                     default_restrictions = {
                         domain: default_restriction_map[integral_type]
@@ -412,12 +414,12 @@ def max_subdomain_ids(self) -> dict[str, int]:
         return _compute_max_subdomain_ids(self.integral_data)
 
     @cached_property
-    def argument_elements(self) -> list[Any]:
+    def argument_elements(self) -> tuple[Any]:
         """The set of elements the arguments in the form."""
         return tuple(f.ufl_element() for f in self.original_form.arguments())
 
     @property
-    def coefficient_elements(self) -> list[Any]:
+    def coefficient_elements(self) -> tuple[Any]:
         """The set of elements used for coefficients in the form."""
         return self._coefficient_elements
 

ufl/domain.py

@@ -16,6 +16,8 @@
     from ufl.core.expr import Expr
     from ufl.finiteelement import AbstractFiniteElement  # To avoid cyclic import when type-hinting.
     from ufl.form import Form
+    from ufl.integral import Integral
+
 from ufl.cell import AbstractCell, CellSequence
 from ufl.core.ufl_id import attach_ufl_id
 from ufl.core.ufl_type import UFLObject
@@ -339,7 +341,7 @@ def as_domain(domain):
         return domain
 
 
-def sort_domains(domains: Sequence[AbstractDomain]) -> tuple[AbstractDomain, ...]:
+def sort_domains(domains: Iterable[AbstractDomain]) -> tuple[AbstractDomain, ...]:
     """Sort domains in a canonical ordering.
 
     Args:
@@ -388,7 +390,7 @@ def join_domains(domains: Sequence[AbstractDomain], expand_mesh_sequence: bool =
 
 
 def extract_domains(
-    expr: Expr | Form,
+    expr: Expr | Form | Integral,
     expand_mesh_sequence: bool = True,
 ) -> tuple[AbstractDomain, ...]:
     """Return all domains expression is defined on.

ufl/form.py

@@ -43,9 +43,12 @@ def _sorted_integrals(integrals: typing.Iterable[Integral]) -> tuple[Integral, .
     """
     # Group integrals in multilevel dict by keys
     # [domain][integral_type][subdomain_id]
-    integrals_dict = defaultdict(
-        lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
-    )
+    from ufl.classes import AbstractDomain
+
+    integrals_dict: dict[
+        AbstractDomain,
+        dict[str, dict[tuple[typing.Any, ...], dict[int | tuple[int, ...], list[Integral]]]],
+    ] = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(list))))
     for integral in integrals:
         d = integral.ufl_domain()
         if d is None:
@@ -314,7 +317,7 @@ def __init__(self, integrals: list[Integral]):
         self._signature = None
 
         # Never use this internally in ufl!
-        self._cache = {}
+        self._cache: dict[typing.Any, typing.Any] = {}
 
     # --- Accessor interface ---