Start fixing 'check_untyped_defs' for 'ffcx.codegeneration.*'

Author: schnellerhase

ffcx/codegeneration/geometry.py

@@ -76,19 +76,19 @@ def reference_facet_volume(tablename, cellname):
     """Write a reference facet volume."""
     celltype = getattr(basix.CellType, cellname)
     volumes = basix.cell.facet_reference_volumes(celltype)
-    for i in volumes[1:]:
-        if not np.isclose(i, volumes[0]):
+    for i in volumes[1:]:  # type: ignore
+        if not np.isclose(i, volumes[0]):  # type: ignore
             raise ValueError("Reference facet volume not supported for this cell type.")
     symbol = L.Symbol(f"{cellname}_{tablename}", dtype=L.DataType.REAL)
-    return L.VariableDecl(symbol, volumes[0])
+    return L.VariableDecl(symbol, volumes[0])  # type: ignore
 
 
 def reference_cell_edge_vectors(tablename, cellname):
     """Write reference edge vectors."""
     celltype = getattr(basix.CellType, cellname)
     topology = basix.topology(celltype)
     geometry = basix.geometry(celltype)
-    edge_vectors = [geometry[j] - geometry[i] for i, j in topology[1]]
+    edge_vectors = [geometry[j] - geometry[i] for i, j in topology[1]]  # type: ignore
     out = np.array(edge_vectors)
     symbol = L.Symbol(f"{cellname}_{tablename}", dtype=L.DataType.REAL)
     return L.ArrayDecl(symbol, values=out, const=True)
@@ -108,10 +108,11 @@ def reference_facet_edge_vectors(tablename, cellname):
     edge_vectors = []
     for facet in topology[-2]:
         if len(facet) == 3:
-            edge_vectors += [geometry[facet[j]] - geometry[facet[i]] for i, j in triangle_edges]
+            edge_vectors += [geometry[facet[j]] - geometry[facet[i]] for i, j in triangle_edges]  # type: ignore
         elif len(facet) == 4:
             edge_vectors += [
-                geometry[facet[j]] - geometry[facet[i]] for i, j in quadrilateral_edges
+                geometry[facet[j]] - geometry[facet[i]]
+                for i, j in quadrilateral_edges  # type: ignore
             ]
         else:
             raise ValueError("Only triangular and quadrilateral faces supported.")

ffcx/codegeneration/integral_generator.py

@@ -209,7 +209,7 @@ def generate_geometry_tables(self):
             ufl.geometry.ReferenceNormal: "reference_normals",
             ufl.geometry.FacetOrientation: "facet_orientation",
         }
-        cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()}  # type: ignore
+        cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()}
 
         for integrand in self.ir.expression.integrand.values():
             for attr in integrand["factorization"].nodes.values():

ffcx/codegeneration/jit.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import importlib
+import importlib.util
 import io
 import logging
 import os
@@ -109,7 +110,7 @@ def get_cached_module(module_name, object_names, cache_dir, timeout):
         for i in range(timeout):
             if os.path.exists(ready_name):
                 spec = finder.find_spec(module_name)
-                if spec is None:
+                if spec is None or spec.loader is None:
                     raise ModuleNotFoundError("Unable to find JIT module.")
                 compiled_module = importlib.util.module_from_spec(spec)
                 spec.loader.exec_module(compiled_module)
@@ -409,11 +410,12 @@ def _load_objects(cache_dir, module_name, object_names):
     # (new) modules are found
     finder.invalidate_caches()
     spec = finder.find_spec(module_name)
-    if spec is None:
+    if spec is None or spec.loader is None:
         raise ModuleNotFoundError("Unable to find JIT module.")
 
     # Load module
     compiled_module = importlib.util.module_from_spec(spec)
+
     spec.loader.exec_module(compiled_module)
 
     compiled_objects = []

ffcx/codegeneration/lnodes.py

@@ -298,6 +298,9 @@ def __eq__(self, other):
 
     def __float__(self):
         """Convert to float."""
+        if isinstance(self.value, complex):
+            raise RuntimeError("__float__ of a complex value not defined.")
+
         return float(self.value)
 
     def __repr__(self):
@@ -316,6 +319,10 @@ def __init__(self, value):
         self.value = value
         self.dtype = DataType.INT
 
+    def __int__(self):
+        """Overwrites int()."""
+        return self.value
+
     def __eq__(self, other):
         """Check equality."""
         return isinstance(other, LiteralInt) and self.value == other.value
@@ -447,6 +454,8 @@ def __eq__(self, other):
 class BinOp(LExprOperator):
     """A binary operator."""
 
+    op = ""
+
     def __init__(self, lhs, rhs):
         """Initialise."""
         self.lhs = as_lexpr(lhs)
@@ -955,7 +964,6 @@ def __eq__(self, other):
         """Check equality."""
         return (
             isinstance(other, type(self))
-            and self.typename == other.typename
             and self.symbol == other.symbol
             and self.value == other.value
         )

ffcx/codegeneration/symbols.py

@@ -185,10 +185,7 @@ def element_table(self, tabledata, entity_type, restriction):
         else:
             entity = self.entity(entity_type, restriction)
 
-        if tabledata.is_piecewise:
-            iq = 0
-        else:
-            iq = self.quadrature_loop_index
+        iq = 0 if tabledata.is_piecewise else self.quadrature_loop_index
 
         if tabledata.is_permuted:
             qp = self.quadrature_permutation[0]

pyproject.toml

@@ -83,9 +83,12 @@ disallow_any_unimported = false # most of these come from UFL
 
 [[tool.mypy.overrides]]
 module = ["ffcx.ir.*", "ffcx.codegeneration.*"]
-check_untyped_defs = false
 disallow_untyped_defs = false 
 
+[[tool.mypy.overrides]]
+module = ["ffcx.ir.*"]
+check_untyped_defs = false
+
 [tool.ruff]
 line-length = 100
 indent-width = 4