Initial mypy support

.github/workflows/pythonapp.yml

@@ -46,6 +46,8 @@ jobs:
           ruff format --check .
       - name: Install UFL
         run: python -m pip install .[ci]
+      - name: Lint with mypy
+        run: mypy -p ufl
       - name: Run unit tests
         run: |
           python -m pytest -n auto --cov=ufl/ --junitxml=junit/test-results-${{ matrix.os }}-${{ matrix.python-version }}.xml test/

pyproject.toml

@@ -24,7 +24,8 @@ issues = "https://github.com/FEniCS/ufl/issues"
 funding = "https://numfocus.org/donate"
 
 [project.optional-dependencies]
-lint = ["ruff"]
+lint = ["mypy", "ruff"]
+typing = ["types-colorama"]
 docs = ["sphinx", "sphinx_rtd_theme"]
 test = ["pytest"]
 ci = [
@@ -35,6 +36,7 @@ ci = [
     "fenics-ufl[docs]",
     "fenics-ufl[lint]",
     "fenics-ufl[test]",
+    "fenics-ufl[typing]",
 ]
 
 [tool.setuptools]
@@ -82,3 +84,13 @@ allowed-confusables = ["𝐚", "𝐀", "∕", "γ", "⨯", "∨"]
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"
+
+[tool.mypy]
+disable_error_code = [
+    "arg-type",
+    "assignment",
+    "attr-defined",
+    "misc",
+    "operator",
+    "type-var",
+]

ufl/algorithms/formsplitter.py

@@ -8,6 +8,7 @@
 #
 # Modified by Cecile Daversin-Catty, 2018
 
+import typing
 from typing import Optional
 
 import numpy as np
@@ -135,7 +136,7 @@ def extract_blocks(
             num_sub_elements = arguments[0].ufl_element().num_sub_elements
             forms = []
             for pi in range(num_sub_elements):
-                form_i = []
+                form_i: typing.List[typing.Optional[object]] = []
                 for pj in range(num_sub_elements):
                     f = fs.split(form, pi, pj)
                     if f.empty():

ufl/algorithms/signature.py

@@ -6,6 +6,7 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
 import hashlib
+import typing
 
 from ufl.algorithms.domain_analysis import canonicalize_metadata
 from ufl.classes import (
@@ -20,6 +21,7 @@
     Label,
     MultiIndex,
 )
+from ufl.core.ufl_type import UFLObject
 from ufl.corealg.traversal import traverse_unique_terminals, unique_post_traversal
 
 
@@ -94,7 +96,7 @@ def compute_terminal_hashdata(expressions, renumbering):
 
 def compute_expression_hashdata(expression, terminal_hashdata) -> bytes:
     """Compute expression hashdata."""
-    cache = {}
+    cache: typing.Dict[UFLObject, bytes] = {}
 
     for expr in unique_post_traversal(expression):
         # Uniquely traverse tree and hash each node

ufl/algorithms/transformer.py

@@ -14,6 +14,7 @@
 # Modified by Anders Logg, 2009-2010
 
 import inspect
+import typing
 
 from ufl.algorithms.map_integrands import map_integrands
 from ufl.classes import Variable, all_ufl_classes
@@ -35,7 +36,7 @@ class Transformer(object):
     transform expression trees from one representation to another.
     """
 
-    _handlers_cache = {}
+    _handlers_cache: typing.Dict[type, typing.Tuple[str, bool]] = {}
 
     def __init__(self, variable_cache=None):
         """Initialise."""

ufl/cell.py

@@ -184,7 +184,7 @@ def peak_types(self) -> typing.Tuple[AbstractCell, ...]:
         return self.sub_entity_types(tdim - 3)
 
 
-_sub_entity_celltypes = {
+_sub_entity_celltypes: typing.Dict[str, list[tuple[str, ...]]] = {
     "vertex": [("vertex",)],
     "interval": [tuple("vertex" for i in range(2)), ("interval",)],
     "triangle": [
@@ -340,7 +340,7 @@ def reconstruct(self, **kwargs: typing.Any) -> Cell:
         return Cell(self._cellname)
 
 
-class TensorProductCell(AbstractCell):
+class TensorProductCell(Cell):
     """Tensor product cell."""
 
     __slots__ = ("_cells", "_tdim")
@@ -358,7 +358,7 @@ def __init__(self, *cells: Cell):
         if not isinstance(self._tdim, numbers.Integral):
             raise ValueError("Expecting integer topological_dimension.")
 
-    def sub_cells(self) -> typing.List[AbstractCell]:
+    def sub_cells(self) -> typing.Tuple[AbstractCell, ...]:
         """Return list of cell factors."""
         return self._cells
 
@@ -398,21 +398,21 @@ def num_sub_entities(self, dim: int) -> int:
     def sub_entities(self, dim: int) -> typing.Tuple[AbstractCell, ...]:
         """Get the sub-entities of the given dimension."""
         if dim < 0 or dim > self._tdim:
-            return []
+            return ()
         if dim == 0:
-            return [Cell("vertex") for i in range(self.num_sub_entities(0))]
+            return tuple(Cell("vertex") for i in range(self.num_sub_entities(0)))
         if dim == self._tdim:
-            return [self]
+            return (self,)
         raise NotImplementedError(f"TensorProductCell.sub_entities({dim}) is not implemented.")
 
     def sub_entity_types(self, dim: int) -> typing.Tuple[AbstractCell, ...]:
         """Get the unique sub-entity types of the given dimension."""
         if dim < 0 or dim > self._tdim:
-            return []
+            return ()
         if dim == 0:
-            return [Cell("vertex")]
+            return (Cell("vertex"),)
         if dim == self._tdim:
-            return [self]
+            return (self,)
         raise NotImplementedError(f"TensorProductCell.sub_entities({dim}) is not implemented.")
 
     def _lt(self, other) -> bool:

ufl/constantvalue.py

@@ -10,6 +10,7 @@
 # Modified by Massimiliano Leoni, 2016.
 
 import numbers
+import typing
 from math import atan2
 
 import ufl
@@ -65,7 +66,7 @@ class Zero(ConstantValue):
 
     __slots__ = ("ufl_free_indices", "ufl_index_dimensions", "ufl_shape")
 
-    _cache = {}
+    _cache: typing.Dict[typing.Tuple[int], "Zero"] = {}
 
     def __getnewargs__(self):
         """Get new args."""
@@ -362,7 +363,7 @@ class IntValue(RealValue):
 
     __slots__ = ()
 
-    _cache = {}
+    _cache: typing.Dict[int, "IntValue"] = {}
 
     def __getnewargs__(self):
         """Get new args."""

ufl/core/expr.py

@@ -16,9 +16,10 @@
 # Modified by Anders Logg, 2008
 # Modified by Massimiliano Leoni, 2016
 
+import typing
 import warnings
 
-from ufl.core.ufl_type import UFLType, update_ufl_type_attributes
+from ufl.core.ufl_type import UFLObject, UFLType, update_ufl_type_attributes
 
 
 class Expr(object, metaclass=UFLType):
@@ -204,10 +205,10 @@ def __init__(self):
 
     # A global dict mapping language_operator_name to the type it
     # produces
-    _ufl_language_operators_ = {}
+    _ufl_language_operators_: typing.Dict[str, object] = {}
 
     # List of all terminal modifier types
-    _ufl_terminal_modifiers_ = []
+    _ufl_terminal_modifiers_: typing.List[UFLObject] = []
 
     # --- Mechanism for profiling object creation and deletion ---
 

ufl/core/multiindex.py

@@ -8,6 +8,8 @@
 #
 # Modified by Massimiliano Leoni, 2016.
 
+import typing
+
 from ufl.core.terminal import Terminal
 from ufl.core.ufl_type import ufl_type
 from ufl.utils.counted import Counted
@@ -30,7 +32,7 @@ class FixedIndex(IndexBase):
 
     __slots__ = ("_hash", "_value")
 
-    _cache = {}
+    _cache: typing.Dict[int, "FixedIndex"] = {}
 
     def __getnewargs__(self):
         """Get new args."""
@@ -116,7 +118,7 @@ class MultiIndex(Terminal):
 
     __slots__ = ("_indices",)
 
-    _cache = {}
+    _cache: typing.Dict[typing.Tuple[int], "MultiIndex"] = {}
 
     def __getnewargs__(self):
         """Get new args."""

ufl/core/ufl_type.py

@@ -416,18 +416,18 @@ class UFLType(type):
     _ufl_handler_name_ = "ufl_type"
 
     # A global array of all Expr and BaseForm subclasses, indexed by typecode
-    _ufl_all_classes_ = []
+    _ufl_all_classes_: typing.List[UFLObject] = []
 
     # A global set of all handler names added
-    _ufl_all_handler_names_ = set()
+    _ufl_all_handler_names_: typing.Set[str] = set()
 
     # A global array of the number of initialized objects for each
     # typecode
-    _ufl_obj_init_counts_ = []
+    _ufl_obj_init_counts_: typing.List[int] = []
 
     # A global array of the number of deleted objects for each
     # typecode
-    _ufl_obj_del_counts_ = []
+    _ufl_obj_del_counts_: typing.List[int] = []
 
     # Type trait: If the type is abstract.  An abstract class cannot
     # be instantiated and does not need all properties specified.

ufl/corealg/multifunction.py

@@ -7,6 +7,7 @@
 #
 # Modified by Massimiliano Leoni, 2016
 
+import typing
 from inspect import signature
 
 from ufl.core.expr import Expr
@@ -46,7 +47,7 @@ class MultiFunction(object):
     algorithm object. Of course Python's function call overhead still applies.
     """
 
-    _handlers_cache = {}
+    _handlers_cache: typing.Dict[type, typing.Tuple[typing.List[str], bool]] = {}
 
     def __init__(self):
         """Initialise."""

ufl/exprequals.py

@@ -1,12 +1,5 @@
 """Expr equals."""
 
-from collections import defaultdict
-
-hash_total = defaultdict(int)
-hash_collisions = defaultdict(int)
-hash_equals = defaultdict(int)
-hash_notequals = defaultdict(int)
-
 
 def expr_equals(self, other):
     """Checks whether the two expressions are represented the exact same way.

ufl/finiteelement.py

@@ -53,7 +53,7 @@ def __hash__(self) -> int:
         """Return a hash."""
 
     @_abc.abstractmethod
-    def __eq__(self, other: AbstractFiniteElement) -> bool:
+    def __eq__(self, other: object) -> bool:
         """Check if this element is equal to another element."""
 
     @_abc.abstractproperty
@@ -112,7 +112,7 @@ def sub_elements(self) -> _typing.List:
         of sub-elements.
         """
 
-    def __ne__(self, other: AbstractFiniteElement) -> bool:
+    def __ne__(self, other: object) -> bool:
         """Check if this element is different to another element."""
         return not self.__eq__(other)
 

ufl/functionspace.py

@@ -80,7 +80,7 @@ def components(self) -> typing.Dict[typing.Tuple[int, ...], int]:
         if len(self._ufl_element.sub_elements) == 0:
             return {(): 0}
 
-        components = {}
+        components: dict[typing.Tuple[int, ...], int] = {}
         offset = 0
         c_offset = 0
         for s in self.ufl_sub_spaces():

ufl/mathfunctions.py

@@ -421,7 +421,7 @@ def evaluate(self, x, mapping, component, index_values):
         """Evaluate."""
         a = self.ufl_operands[1].evaluate(x, mapping, component, index_values)
         try:
-            import scipy.special
+            import scipy.special  # type: ignore
         except ImportError:
             raise ValueError(
                 "You must have scipy installed to evaluate bessel functions in python."

ufl/objects.py

@@ -20,10 +20,10 @@
     globals()[measure_name] = Measure(integral_type)
 
 # TODO: Firedrake hack, remove later
-ds_tb = ds_b + ds_t  # noqa: F821
+ds_tb = ds_b + ds_t  # type: ignore # noqa: F821
 
 # Default measure dX including both uncut and cut cells
-dX = dx + dC  # noqa: F821
+dX = dx + dC  # type: ignore # noqa: F821
 
 # Create objects for builtin known cell types
 vertex = Cell("vertex")

ufl/pullback.py

@@ -462,7 +462,7 @@ class SymmetricPullback(AbstractPullback):
     """Pull back for an element with symmetry."""
 
     def __init__(
-        self, element: _AbstractFiniteElement, symmetry: typing.Dict[typing.tuple[int, ...], int]
+        self, element: _AbstractFiniteElement, symmetry: typing.Dict[typing.Tuple[int, ...], int]
     ):
         """Initalise.