Add typing generics to CoordinateElement

Author: schnellerhase

python/dolfinx/fem/element.py

@@ -6,6 +6,7 @@
 """Finite elements."""
 
 from functools import singledispatch
+from typing import Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -14,8 +15,10 @@
 import basix.ufl
 from dolfinx import cpp as _cpp
 
+_T_CE = TypeVar("_T_CE", np.float32, np.float64)
 
-class CoordinateElement:
+
+class CoordinateElement(Generic[_T_CE]):
     """Coordinate element describing the geometry map for mesh cells."""
 
     _cpp_object: _cpp.fem.CoordinateElement_float32 | _cpp.fem.CoordinateElement_float64
@@ -61,10 +64,8 @@ def create_dof_layout(self) -> _cpp.fem.ElementDofLayout:
         return self._cpp_object.create_dof_layout()
 
     def push_forward(
-        self,
-        X: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-        cell_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
+        self, X: npt.NDArray[_T_CE], cell_geometry: npt.NDArray[_T_CE]
+    ) -> npt.NDArray[_T_CE]:
         """Push points on the reference cell forward to the physical cell.
 
         Args:
@@ -81,10 +82,8 @@ def push_forward(
         return self._cpp_object.push_forward(X, cell_geometry)
 
     def pull_back(
-        self,
-        x: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-        cell_geometry: npt.NDArray[np.float32] | npt.NDArray[np.float64],
-    ) -> npt.NDArray[np.float32] | npt.NDArray[np.float64]:
+        self, x: npt.NDArray[_T_CE], cell_geometry: npt.NDArray[_T_CE]
+    ) -> npt.NDArray[_T_CE]:
         """Pull points on the physical cell back to the reference cell.
 
         For non-affine cells, the pull-back is a nonlinear operation.
@@ -124,7 +123,7 @@ def coordinate_element(
     degree: int,
     variant=int(basix.LagrangeVariant.unset),
     dtype: npt.DTypeLike = np.float64,
-):
+) -> CoordinateElement:
     """Create a Lagrange CoordinateElement from element metadata.
 
     Coordinate elements are typically used to create meshes.
@@ -147,7 +146,7 @@ def coordinate_element(
 
 
 @coordinate_element.register(basix.finite_element.FiniteElement)
-def _(e: basix.finite_element.FiniteElement):
+def _(e: basix.finite_element.FiniteElement) -> CoordinateElement:
     """Create a Lagrange CoordinateElement from a Basix finite element.
 
     Coordinate elements are typically used when creating meshes.