Add typing generics Function

Author: schnellerhase

python/dolfinx/fem/function.py

@@ -311,7 +311,7 @@ def dtype(self) -> npt.DTypeLike:
         return np.dtype(self._cpp_object.dtype)
 
 
-class Function(ufl.Coefficient):
+class Function(ufl.Coefficient, Generic[_S]):
     """A finite element function.
 
     A finite element function is represented by a function space
@@ -327,10 +327,12 @@ class Function(ufl.Coefficient):
         | _cpp.fem.Function_float64
     )
 
+    _x: la.Vector[_S]
+
     def __init__(
         self,
         V: FunctionSpace,
-        x: la.Vector | None = None,
+        x: la.Vector[_S] | None = None,
         name: str | None = None,
         dtype: npt.DTypeLike | None = None,
     ):
@@ -395,7 +397,7 @@ def function_space(self) -> FunctionSpace:
         """FunctionSpace that the Function is defined on."""
         return self._V
 
-    def eval(self, x: npt.ArrayLike, cells: npt.ArrayLike, u=None) -> np.ndarray:
+    def eval(self, x: npt.ArrayLike, cells: npt.NDArray[np.int32], u=None) -> npt.NDArray[_S]:
         """Evaluate Function at points x.
 
         Points where x has shape (num_points, 3), and cells has shape
@@ -431,7 +433,7 @@ def eval(self, x: npt.ArrayLike, cells: npt.ArrayLike, u=None) -> np.ndarray:
         return u
 
     def interpolate_nonmatching(
-        self, u0: Function, cells: npt.NDArray[np.int32], interpolation_data: PointOwnershipData
+        self, u0: Function[_S], cells: npt.NDArray[np.int32], interpolation_data: PointOwnershipData
     ) -> None:
         """Interpolate a Function on a non-matching mesh.
 
@@ -447,9 +449,9 @@ def interpolate_nonmatching(
 
     def interpolate(
         self,
-        u0: Callable | Expression | Function,
-        cells0: np.ndarray | None = None,
-        cells1: np.ndarray | None = None,
+        u0: Callable | Expression[_S] | Function[_S],
+        cells0: npt.NDArray[np.int32] | None = None,
+        cells1: npt.NDArray[np.int32] | None = None,
     ) -> None:
         """Interpolate an expression.
 
@@ -495,7 +497,7 @@ def _(e0: Expression):
             )
             self._cpp_object.interpolate_f(np.asarray(u0(x), dtype=self.dtype), cells0)
 
-    def copy(self) -> Function:
+    def copy(self) -> Function[_S]:
         """Create a copy of the Function.
 
         The function space is shared and the degree-of-freedom vector is
@@ -511,12 +513,12 @@ def copy(self) -> Function:
         )
 
     @property
-    def x(self) -> la.Vector:
+    def x(self) -> la.Vector[_S]:
         """Vector holding the degrees-of-freedom."""
         return self._x
 
     @property
-    def dtype(self) -> np.dtype:
+    def dtype(self) -> npt.DTypeLike:
         """Function value dtype."""
         return np.dtype(self._cpp_object.x.array.dtype)
 
@@ -529,11 +531,11 @@ def name(self) -> str:
     def name(self, name):
         self._cpp_object.name = name
 
-    def __str__(self):
+    def __str__(self) -> str:
         """Pretty print representation."""
         return self.name
 
-    def sub(self, i: int) -> Function:
+    def sub(self, i: int) -> Function[_S]:
         """Return a sub-function (a view into the ``Function``).
 
         Sub-functions are indexed ``i = 0, ..., N-1``, where ``N`` is
@@ -552,7 +554,7 @@ def sub(self, i: int) -> Function:
         """
         return Function(self._V.sub(i), self.x, name=f"{self!s}_{i}")
 
-    def split(self) -> tuple[Function, ...]:
+    def split(self) -> tuple[Function[_S], ...]:
         """Extract (any) sub-functions.
 
         A sub-function can be extracted from a discrete function that is
@@ -567,7 +569,7 @@ def split(self) -> tuple[Function, ...]:
             raise RuntimeError("No subfunctions to extract")
         return tuple(self.sub(i) for i in range(num_sub_spaces))
 
-    def collapse(self) -> Function:
+    def collapse(self) -> Function[_S]:
         """Create a collapsed version of this Function."""
         u_collapsed = self._cpp_object.collapse()  # type: ignore
         V_collapsed = FunctionSpace(