Add typing generics dirichletbc

Author: schnellerhase

python/dolfinx/fem/bcs.py

@@ -14,6 +14,7 @@
 
 import numbers
 from collections.abc import Callable, Iterable
+from typing import Generic, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -91,7 +92,10 @@ def locate_dofs_topological(
     return _cpp.fem.locate_dofs_topological(_V, entity_dim, _entities, remote)
 
 
-class DirichletBC:
+_T = TypeVar("_T", np.float32, np.float64, np.complex64, np.complex128)
+
+
+class DirichletBC(Generic[_T]):
     """Representation of Dirichlet boundary conditions.
 
     The conditions are imposed on a linear system.
@@ -120,18 +124,17 @@ class initialiser. This class is combined with different
         self._cpp_object = bc
 
     @property
-    def g(self) -> Function | Constant | np.ndarray:
+    def g(self) -> Function | Constant:
         """The boundary condition value(s)."""
+        # TODO: needs to be wrapped into Function or Constant
         return self._cpp_object.value
 
     @property
     def function_space(self) -> dolfinx.fem.FunctionSpace:
         """Function space on which the boundary condition is defined."""
         return self._cpp_object.function_space
 
-    def set(
-        self, x: npt.NDArray, x0: npt.NDArray[np.int32] | None = None, alpha: float = 1
-    ) -> None:
+    def set(self, x: npt.NDArray[_T], x0: npt.NDArray[_T] | None = None, alpha: float = 1) -> None:
         """Set array entries that are constrained by a Dirichlet condition.
 
         Entries in ``x`` that are constrained by a Dirichlet boundary
@@ -172,10 +175,10 @@ def dof_indices(self) -> tuple[npt.NDArray[np.int32], int]:
 
 
 def dirichletbc(
-    value: Function | Constant | np.ndarray,
+    value: Function | Constant | npt.NDArray[_T],
     dofs: npt.NDArray[np.int32],
     V: dolfinx.fem.FunctionSpace | None = None,
-) -> DirichletBC:
+) -> DirichletBC[_T]:
     """Representation of Dirichlet boundary condition.
 
     Args:
@@ -232,8 +235,8 @@ def dirichletbc(
 
 
 def bcs_by_block(
-    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC]
-) -> list[list[DirichletBC]]:
+    spaces: Iterable[FunctionSpace | None], bcs: Iterable[DirichletBC[_T]]
+) -> list[list[DirichletBC[_T]]]:
     """Arrange boundary conditions by the space that they constrain.
 
     Given a sequence of function spaces ``spaces`` and a sequence of