More typing

.github/workflows/fenicsx-refs.env

@@ -3,4 +3,4 @@ ufl_ref=main
 ffcx_repository=FEniCS/ffcx
 ffcx_ref=main
 dolfinx_repository=FEniCS/dolfinx
-dolfinx_ref=main
+dolfinx_ref=schnellerhase/mypy-demo

python/basix/ufl.py

@@ -200,7 +200,9 @@ def sub_elements(self) -> _typing.Sequence[_AbstractFiniteElement]:
 
     # Basix specific functions
     @_abstractmethod
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(
+        self, nderivs: int, points: _npt.NDArray[np.floating]
+    ) -> _npt.NDArray[np.floating]:
         """Tabulate the basis functions of the element.
 
         Args:
@@ -453,7 +455,9 @@ def __hash__(self) -> int:
     def basix_hash(self) -> int | None:
         return self._element.hash()
 
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(
+        self, nderivs: int, points: _npt.NDArray[np.floating]
+    ) -> _npt.NDArray[np.floating]:
         tab = self._element.tabulate(nderivs, points)
         # TODO: update FFCx to remove the need for transposing here
         return tab.transpose((0, 1, 3, 2)).reshape((tab.shape[0], tab.shape[1], -1))  # type: ignore
@@ -601,7 +605,7 @@ def __eq__(self, other) -> bool:
     def __hash__(self) -> int:
         return super().__hash__()
 
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.NDArray[np.float64]:
         tables = self._element.tabulate(nderivs, points)
         output = []
         for tbl in tables:  # type: ignore
@@ -769,7 +773,7 @@ def is_mixed(self) -> bool:
     def degree(self) -> int:
         return max((e.degree for e in self._sub_elements), default=-1)
 
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.NDArray[np.float64]:
         tables = []
         results = [e.tabulate(nderivs, points) for e in self._sub_elements]
         for deriv_tables in zip(*results):
@@ -1040,7 +1044,7 @@ def is_quadrature(self) -> bool:
     def is_real(self) -> bool:
         return self._sub_element.is_real
 
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.NDArray[np.float64]:
         assert len(self._block_shape) == 1  # TODO: block shape
         assert self.reference_value_size == self._block_size  # TODO: remove this assumption
         output = []
@@ -1277,7 +1281,9 @@ def __eq__(self, other) -> bool:
     def __hash__(self) -> int:
         return super().__hash__()
 
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(
+        self, nderivs: int, points: _npt.NDArray[np.floating]
+    ) -> _npt.NDArray[np.floating]:
         if nderivs > 0:
             raise ValueError("Cannot take derivatives of Quadrature element.")
 
@@ -1420,7 +1426,9 @@ def __hash__(self) -> int:
     def dtype(self) -> _npt.DTypeLike:
         raise NotImplementedError()
 
-    def tabulate(self, nderivs: int, points: _npt.NDArray[np.floating]) -> _npt.ArrayLike:
+    def tabulate(
+        self, nderivs: int, points: _npt.NDArray[np.floating]
+    ) -> _npt.NDArray[np.floating]:
         out = np.zeros((nderivs + 1, len(points), self.reference_value_size**2))
         for v in range(self.reference_value_size):
             out[0, :, self.reference_value_size * v + v] = 1.0
@@ -1574,7 +1582,7 @@ def element(
     symmetry: bool | None = None,
     dof_ordering: list[int] | None = None,
     dtype: _npt.DTypeLike | None = None,
-) -> _ElementBase:
+) -> _BasixElement | _BlockedElement:
     """Create a UFL compatible element using Basix.
 
     Args:
@@ -1915,7 +1923,7 @@ def blocked_element(
     sub_element: _ElementBase,
     shape: tuple[int, ...],
     symmetry: bool | None = None,
-) -> _ElementBase:
+) -> _BlockedElement:
     """Create a UFL compatible blocked element.
 
     Args: