Extend typing checks to demos

.github/workflows/ccpp.yml

@@ -138,7 +138,7 @@ jobs:
           mypy --config-file python/pyproject.toml -p dolfinx # package mode, type checks installed dolfinx (working around name collision in python/ directory)
           cd python/
           mypy test
-          #mypy demo # TODO: Enable this in future
+          mypy demo
 
       - name: Install gmsh and pyvista (and dependencies)
         run: |

.github/workflows/fenicsx-refs.env

@@ -1,5 +1,5 @@
 basix_repository=FEniCS/basix
-basix_ref=main
+basix_ref=schnellerhase/mypy-dolfinx
 ufl_repository=FEniCS/ufl
 ufl_ref=main
 ffcx_repository=FEniCS/ffcx

python/demo/demo_hdg.py

@@ -98,8 +98,8 @@ def u_e(x):
 
 # Create the mesh
 
-n = 8  # Number of elements in each direction
-msh = mesh.create_unit_cube(comm, n, n, n, ghost_mode=mesh.GhostMode.none)
+n_el = 8  # Number of elements in each direction
+msh = mesh.create_unit_cube(comm, n_el, n_el, n_el, ghost_mode=mesh.GhostMode.none)
 
 # We need to create a broken Lagrange space defined over the facets of
 # the mesh. To do so, we require a sub-mesh of the all facets. We begin

python/demo/demo_interpolation-io.py

@@ -119,7 +119,7 @@
     pl.add_text("y-component", font_size=12, color="black", position="upper_edge")
     pl.add_mesh(grid.copy(), component=1, show_edges=True)
 
-    pl.view_xy()
+    pl.view_xy()  # type: ignore
     pl.link_views()
 
     # If pyvista environment variable is set to off-screen (static)

python/demo/demo_lagrange-variants.py

@@ -65,8 +65,8 @@
 if MPI.COMM_WORLD.size == 1:
     for i in range(values.shape[1]):
         plt.plot(lattice, values[:, i])
-    plt.plot(element._element.points, [0] * 11, "ko")
-    plt.ylim([-1, 6])
+    plt.plot(element.basix_element.points, [0] * 11, "ko")
+    plt.ylim((-1, 6))
     plt.savefig("demo_lagrange_variants_equispaced_10.png")
     plt.clf()
 # -
@@ -97,8 +97,8 @@
 if MPI.COMM_WORLD.size == 1:  # Skip this plotting in parallel
     for i in range(values.shape[1]):
         plt.plot(lattice, values[:, i])
-    plt.plot(element._element.points, [0] * 11, "ko")
-    plt.ylim([-1, 6])
+    plt.plot(element.basix_element.points, [0] * 11, "ko")
+    plt.ylim((-1, 6))
     plt.savefig("demo_lagrange_variants_gll_10.png")
     plt.clf()
 # -
@@ -152,7 +152,7 @@ def saw_tooth(x):
         plt.plot(pts, [saw_tooth(i[0]) for i in pts], "k--")
         plt.plot(pts, values, "r-")
         plt.legend(["function", "approximation"])
-        plt.ylim([-0.1, 0.4])
+        plt.ylim((-0.1, 0.4))
         plt.title(variant.name)
         plt.savefig(f"demo_lagrange_variants_interpolation_{variant.name}.png")
         plt.clf()

python/demo/demo_mixed-topology.py

@@ -193,9 +193,9 @@ def marker(x):
 A_scipy = A.to_scipy()
 b_scipy = b.array
 
-x = spsolve(A_scipy, b_scipy)
+x_scipy = spsolve(A_scipy, b_scipy)
 
-print(f"Solution vector norm {np.linalg.norm(x)}")
+print(f"Solution vector norm {np.linalg.norm(x_scipy)}")
 
 # Mixed-topology I/O
 # We manually build a ASCII XDMF file to store the mesh
@@ -235,8 +235,8 @@ def marker(x):
           </DataItem>
         </Geometry>
         <Attribute Name="u" Center="Node" NumberType="float" Precision="8">
-          <DataItem Dimensions="{len(x)}" Format="XML">
-            {" ".join(str(val) for val in x)}
+          <DataItem Dimensions="{len(x_scipy)}" Format="XML">
+            {" ".join(str(val) for val in x_scipy)}
           </DataItem>
        </Attribute>
       </Grid>"""

python/demo/demo_navier-stokes.py

@@ -242,7 +242,7 @@ def f_expr(x):
 
 # We define some simulation parameters
 
-n = 16
+n_el = 16
 num_time_steps = 25
 t_end = 10
 Re = 25  # Reynolds Number
@@ -254,7 +254,7 @@ def f_expr(x):
 # interpolate into for artifact free visualisation.
 
 # +
-msh = mesh.create_unit_square(MPI.COMM_WORLD, n, n)
+msh = mesh.create_unit_square(MPI.COMM_WORLD, n_el, n_el)
 
 # Function spaces for the velocity and for the pressure
 V = fem.functionspace(msh, ("Raviart-Thomas", k + 1))
@@ -413,7 +413,7 @@ def jump(phi, n):
 # We perform the time-stepping as a for-loop
 
 # +
-for n in range(num_time_steps):
+for _ in range(num_time_steps):
     t += delta_t.value
 
     navier_stokes_problem.solve()

python/dolfinx/fem/forms.py

@@ -476,10 +476,23 @@
     return _create_form(form)
 
 
+@typing.overload
+def extract_function_spaces(forms: Form, index: int = 0) -> FunctionSpace | None: ...
+
+
+@typing.overload
 def extract_function_spaces(
-    forms: Form | Sequence[Form] | Sequence[Sequence[Form]],
-    index: int = 0,
-) -> FunctionSpace | list[None | FunctionSpace]:
+    forms: Sequence[Form], index: int = 0
+) -> list[FunctionSpace | None]: ...
+
+
+@typing.overload
+def extract_function_spaces(
+    forms: Sequence[Sequence[Form]], index: int = 0
+) -> list[list[FunctionSpace | None]]: ...
+
+
+def extract_function_spaces(forms, index: int = 0):
     """Extract common function spaces from an array of forms.
 
     If ``forms`` is a list of linear forms, this function returns of list

python/dolfinx/fem/petsc.py

@@ -468,10 +468,10 @@ def _(
                     "Cannot have a entire {'row' if index == 0 else 'column'} of a full of None"
                 )
         is0 = _cpp.la.petsc.create_index_sets(
-            [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[0]]  # type: ignore[union-attr]
+            [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[0]]  # type: ignore
         )
         is1 = _cpp.la.petsc.create_index_sets(
-            [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[1]]  # type: ignore[union-attr]
+            [(Vsub.dofmaps[0].index_map, Vsub.dofmaps[0].index_map_bs) for Vsub in V[1]]  # type: ignore
         )
 
         _bcs = [bc._cpp_object for bc in bcs] if bcs is not None else []

python/dolfinx/io/utils.py

@@ -213,6 +213,7 @@ def read_mesh(
 
         # Get coordinate element, special handling for second order
         # serendipity.
+        basix_el: basix.ufl._BasixElement | basix.ufl._BlockedElement
         num_nodes_per_cell = cells.shape[1]
         if (cell_shape == CellType.quadrilateral and num_nodes_per_cell == 8) or (
             cell_shape == CellType.hexahedron and num_nodes_per_cell == 20

python/dolfinx/plot.py

@@ -6,8 +6,10 @@
 """Support functions for plotting."""
 
 import functools
+from typing import overload
 
 import numpy as np
+import numpy.typing as npt
 
 from dolfinx import cpp as _cpp
 from dolfinx import fem, mesh
@@ -29,8 +31,18 @@
 }
 
 
+@overload
+def vtk_mesh(
+    msh: mesh.Mesh, dim: int | None = None, entities: npt.NDArray[np.int32] | None = None
+): ...
+
+
+@overload
+def vtk_mesh(V: fem.FunctionSpace, entities: npt.NDArray[np.int32] | None = None): ...
+
+
 @functools.singledispatch
-def vtk_mesh(msh: mesh.Mesh, dim: int | None = None, entities=None):
+def vtk_mesh(msh: mesh.Mesh, dim: int | None = None, entities: npt.NDArray[np.int32] | None = None):
     """Create VTK mesh topology data for mesh entities.
 
     The vertex indices in the returned topology array are the indices
@@ -74,8 +86,8 @@
     return topology.reshape(-1), cell_types, msh.geometry.x
 
 
-@vtk_mesh.register
-def _(V: fem.FunctionSpace, entities=None):  # type: ignore
+@vtk_mesh.register  # type: ignore[attr-defined]
+def _(V: fem.FunctionSpace, entities: npt.NDArray[np.int32] | None = None):  # type: ignore
     """Create VTK mesh topology based on the degree-of-freedom coordinates.
 
     This function supports visualisation when the degree of the finite
@@ -113,7 +125,7 @@
     msh = V.mesh
     tdim = msh.topology.dim
     if entities is None:
-        entities = range(msh.topology.index_map(tdim).size_local)
+        entities = np.arange(msh.topology.index_map(tdim).size_local, dtype=np.int32)
 
     dofmap = V.dofmap
     num_dofs_per_cell = V.dofmap.dof_layout.num_dofs
@@ -123,9 +135,9 @@
     vtk_type = (
         _first_order_vtk[cell_type] if degree == 1 else _cpp.io.get_vtk_cell_type(cell_type, tdim)
     )
-    cell_types = np.full(len(entities), vtk_type)
+    cell_types = np.full(entities.size, vtk_type)
 
-    topology = np.zeros((len(entities), num_dofs_per_cell + 1), dtype=np.int32)
+    topology = np.zeros((entities.size, num_dofs_per_cell + 1), dtype=np.int32)
     topology[:, 0] = num_dofs_per_cell
     dofmap_ = dofmap.list
     topology[:, 1:] = dofmap_[entities][:, perm]