From 69820ebfbc9e4b14c61ea830343f3e96174a2ba1 Mon Sep 17 00:00:00 2001
From: schnellerhase <56360279+schnellerhase@users.noreply.github.com>
Date: Wed, 14 Aug 2024 17:35:27 +0200
Subject: [PATCH 1/3] Set mypy no_implicit_optional=true

---
 python/pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/pyproject.toml b/python/pyproject.toml
index baed60b2948..ef5a514e36f 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -73,7 +73,7 @@ markers = [
 # Goal would be to make all of the below True long-term
 disallow_untyped_defs = false
 disallow_any_unimported = false
-no_implicit_optional = false
+no_implicit_optional = true
 check_untyped_defs = false
 warn_return_any = false
 warn_unused_ignores = false

From 5a3d9efba9ec3c738e906336a1b4efe5b70eff89 Mon Sep 17 00:00:00 2001
From: schnellerhase <56360279+schnellerhase@users.noreply.github.com>
Date: Wed, 14 Aug 2024 17:49:03 +0200
Subject: [PATCH 2/3] DROP ME

---
 python/dolfinx/common.py    | 21 +++++++++++----------
 python/dolfinx/fem/forms.py | 19 ++++++++-----------
 python/dolfinx/pkgconfig.py |  2 +-
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/python/dolfinx/common.py b/python/dolfinx/common.py
index be70ce239ab..24e8607a892 100644
--- a/python/dolfinx/common.py
+++ b/python/dolfinx/common.py
@@ -8,7 +8,9 @@
 import functools
 import typing
 
-from dolfinx import cpp as _cpp
+from dolfinx.cpp import Reduction as _Reduction
+from dolfinx.cpp import Timer as _Timer
+from dolfinx.cpp import TimingTyp as _TimingType
 from dolfinx.cpp.common import (
     IndexMap,
     git_commit_hash,
@@ -18,6 +20,8 @@
     has_parmetis,
     has_petsc,
 )
+from dolfinx.cpp.common import list_timings as _list_timings
+from dolfinx.cpp.common import timing as _timing
 
 __all__ = [
     "IndexMap",
@@ -31,12 +35,12 @@
     "has_parmetis",
 ]
 
-TimingType = _cpp.common.TimingType
-Reduction = _cpp.common.Reduction
+TimingType = _TimingType
+Reduction = _Reduction
 
 
 def timing(task: str):
-    return _cpp.common.timing(task)
+    return _timing(task)
 
 
 def list_timings(comm, timing_types: list, reduction=Reduction.max):
@@ -44,7 +48,7 @@ def list_timings(comm, timing_types: list, reduction=Reduction.max):
     wall time, system time or user time. When used in parallel, a
     reduction is applied across all processes. By default, the maximum
     time is shown."""
-    _cpp.common.list_timings(comm, timing_types, reduction)
+    _list_timings(comm, timing_types, reduction)
 
 
 class Timer:
@@ -81,13 +85,10 @@ class Timer:
         list_timings(comm, [TimingType.wall, TimingType.user])
     """
 
-    _cpp_object: _cpp.common.Timer
+    _cpp_object: _Timer
 
     def __init__(self, name: typing.Optional[str] = None):
-        if name is None:
-            self._cpp_object = _cpp.common.Timer()
-        else:
-            self._cpp_object = _cpp.common.Timer(name)
+        self._cpp_object = _Timer(None if name is None else name)
 
     def __enter__(self):
         self._cpp_object.start()
diff --git a/python/dolfinx/fem/forms.py b/python/dolfinx/fem/forms.py
index 6c4ed21311c..6a33fd301fa 100644
--- a/python/dolfinx/fem/forms.py
+++ b/python/dolfinx/fem/forms.py
@@ -254,25 +254,22 @@ def _form(form):
         constants = [c._cpp_object for c in form.constants()]
 
         # Make map from integral_type to subdomain id
-        subdomain_ids = {type: [] for type in sd.get(domain).keys()}
+        subdomain_ids: dict[str, list[int]] = {type: [] for type in sd.get(domain).keys()}
         for integral in form.integrals():
             if integral.subdomain_data() is not None:
                 # Subdomain ids can be strings, its or tuples with strings and ints
                 if integral.subdomain_id() != "everywhere":
                     try:
-                        ids = [sid for sid in integral.subdomain_id() if sid != "everywhere"]
+                        subdomain_ids[integral.integral_type()] += [
+                            sid for sid in integral.subdomain_id() if sid != "everywhere"
+                        ]
                     except TypeError:
                         # If not tuple, but single integer id
-                        ids = [integral.subdomain_id()]
-                else:
-                    ids = []
-                subdomain_ids[integral.integral_type()].append(ids)
+                        subdomain_ids[integral.integral_type()].append(integral.subdomain_id())
 
-        # Chain and sort subdomain ids
-        for itg_type, marker_ids in subdomain_ids.items():
-            flattened_ids = list(chain.from_iterable(marker_ids))
-            flattened_ids.sort()
-            subdomain_ids[itg_type] = flattened_ids
+        # Sort subdomain ids
+        for val in subdomain_ids.values():
+            val.sort()
 
         # Subdomain markers (possibly empty list for some integral types)
         subdomains = {
diff --git a/python/dolfinx/pkgconfig.py b/python/dolfinx/pkgconfig.py
index 41150a9fbef..6d4f9e0fe4e 100644
--- a/python/dolfinx/pkgconfig.py
+++ b/python/dolfinx/pkgconfig.py
@@ -37,7 +37,7 @@ def parse(package):
         "-l": "libraries",
     }
 
-    result = {x: [] for x in parse_map.values()}
+    result: dict[str, list[str]] = {x: [] for x in parse_map.values()}
 
     # Execute the query to pkg-config and clean the result
     out = _pkgconfig_query(package + " --cflags --libs")[1]

From 96a8a7dba871f4c43b734169be0290219d562573 Mon Sep 17 00:00:00 2001
From: schnellerhase <56360279+schnellerhase@users.noreply.github.com>
Date: Wed, 14 Aug 2024 23:36:35 +0200
Subject: [PATCH 3/3] Set mypy disallow_untyped_defs=true and adapt

---
 python/dolfinx/__init__.py          |   2 +-
 python/dolfinx/common.py            |  34 ++++++---
 python/dolfinx/fem/bcs.py           |  20 +++--
 python/dolfinx/fem/dofmap.py        |  17 +++--
 python/dolfinx/fem/element.py       |   6 +-
 python/dolfinx/fem/function.py      | 111 ++++++++++++++++------------
 python/dolfinx/geometry.py          |  28 +++++--
 python/dolfinx/graph.py             |   7 +-
 python/dolfinx/io/utils.py          |  65 +++++++++++-----
 python/dolfinx/jit.py               |  20 +++--
 python/dolfinx/la.py                |  30 ++++----
 python/dolfinx/mesh.py              |  84 +++++++++++++--------
 python/dolfinx/nls/petsc.py         |   4 +-
 python/dolfinx/pkgconfig.py         |   6 +-
 python/dolfinx/plot.py              |  11 ++-
 python/dolfinx/wrappers/__init__.py |   8 +-
 python/pyproject.toml               |   2 +-
 17 files changed, 288 insertions(+), 167 deletions(-)

diff --git a/python/dolfinx/__init__.py b/python/dolfinx/__init__.py
index dcdfa35aa3f..90d74052a58 100644
--- a/python/dolfinx/__init__.py
+++ b/python/dolfinx/__init__.py
@@ -43,7 +43,7 @@
 del _cpp, sys
 
 
-def get_include(user=False):
+def get_include(user: bool = False) -> str:
     import os
 
     d = os.path.dirname(__file__)
diff --git a/python/dolfinx/common.py b/python/dolfinx/common.py
index 24e8607a892..341dc2f51bb 100644
--- a/python/dolfinx/common.py
+++ b/python/dolfinx/common.py
@@ -7,6 +7,9 @@
 
 import functools
 import typing
+from types import TracebackType
+
+from mpi4py import MPI as _MPI
 
 from dolfinx.cpp import Reduction as _Reduction
 from dolfinx.cpp import Timer as _Timer
@@ -39,11 +42,13 @@
 Reduction = _Reduction
 
 
-def timing(task: str):
+def timing(task: str) -> tuple[int, float, float, float]:
     return _timing(task)
 
 
-def list_timings(comm, timing_types: list, reduction=Reduction.max):
+def list_timings(
+    comm: _MPI.Comm, timing_types: list, reduction: _Reduction = Reduction.max
+) -> None:
     """Print out a summary of all Timer measurements, with a choice of
     wall time, system time or user time. When used in parallel, a
     reduction is applied across all processes. By default, the maximum
@@ -87,35 +92,40 @@ class Timer:
 
     _cpp_object: _Timer
 
-    def __init__(self, name: typing.Optional[str] = None):
+    def __init__(self, name: typing.Optional[str] = None) -> None:
         self._cpp_object = _Timer(None if name is None else name)
 
-    def __enter__(self):
+    def __enter__(self) -> typing.Self:
         self._cpp_object.start()
         return self
 
-    def __exit__(self, *args):
+    def __exit__(
+        self,
+        exception_type: typing.Optional[BaseException],
+        exception_value: typing.Optional[BaseException],
+        traceback: typing.Optional[TracebackType],
+    ) -> None:
         self._cpp_object.stop()
 
-    def start(self):
+    def start(self) -> None:
         self._cpp_object.start()
 
-    def stop(self):
+    def stop(self) -> float:
         return self._cpp_object.stop()
 
-    def resume(self):
+    def resume(self) -> None:
         self._cpp_object.resume()
 
-    def elapsed(self):
+    def elapsed(self) -> float:
         return self._cpp_object.elapsed()
 
 
-def timed(task: str):
+def timed(task: str) -> typing.Callable[..., typing.Any]:
     """Decorator for timing functions."""
 
-    def decorator(func):
+    def decorator(func: typing.Callable[..., typing.Any]) -> typing.Callable[..., typing.Any]:
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
             with Timer(task):
                 return func(*args, **kwargs)
 
diff --git a/python/dolfinx/fem/bcs.py b/python/dolfinx/fem/bcs.py
index 1972740bfcb..5520d30cb3c 100644
--- a/python/dolfinx/fem/bcs.py
+++ b/python/dolfinx/fem/bcs.py
@@ -98,7 +98,15 @@ class DirichletBC:
         _cpp.fem.DirichletBC_float64,
     ]
 
-    def __init__(self, bc):
+    def __init__(
+        self,
+        bc: typing.Union[
+            _cpp.fem.DirichletBC_complex64,
+            _cpp.fem.DirichletBC_complex128,
+            _cpp.fem.DirichletBC_float32,
+            _cpp.fem.DirichletBC_float64,
+        ],
+    ):
         """Representation of Dirichlet boundary condition which is imposed on
         a linear system.
 
@@ -195,8 +203,8 @@ def dirichletbc(
 
 
 def bcs_by_block(
-    spaces: typing.Iterable[typing.Union[dolfinx.fem.FunctionSpace, None]],
-    bcs: typing.Iterable[DirichletBC],
+    spaces: list[typing.Union[dolfinx.fem.FunctionSpace, None]],
+    bcs: list[DirichletBC],
 ) -> list[list[DirichletBC]]:
     """Arrange Dirichlet boundary conditions by the function space that
     they constrain.
@@ -207,8 +215,10 @@ def bcs_by_block(
     ``space[i]``.
     """
 
-    def _bc_space(V, bcs):
+    def _bc_space(
+        V: list[typing.Union[dolfinx.fem.FunctionSpace, None]], bcs: list[DirichletBC]
+    ) -> list[DirichletBC]:
         """Return list of bcs that have the same space as V"""
-        return [bc for bc in bcs if V.contains(bc.function_space)]
+        return [bc for bc in bcs if bc.function_space in V]
 
     return [_bc_space(V, bcs) if V is not None else [] for V in spaces]
diff --git a/python/dolfinx/fem/dofmap.py b/python/dolfinx/fem/dofmap.py
index 0355a2cd860..059c694a742 100644
--- a/python/dolfinx/fem/dofmap.py
+++ b/python/dolfinx/fem/dofmap.py
@@ -4,7 +4,12 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
+import numpy as np
+import numpy.typing as npt
+
 from dolfinx import cpp as _cpp
+from dolfinx.common import IndexMap
+from dolfinx.fem import ElementDofLayout
 
 
 class DofMap:
@@ -19,7 +24,7 @@ class DofMap:
     def __init__(self, dofmap: _cpp.fem.DofMap):
         self._cpp_object = dofmap
 
-    def cell_dofs(self, cell_index: int):
+    def cell_dofs(self, cell_index: int) -> npt.NDArray[np.int32]:
         """Cell local-global dof map
 
         Args:
@@ -31,26 +36,26 @@ def cell_dofs(self, cell_index: int):
         return self._cpp_object.cell_dofs(cell_index)
 
     @property
-    def bs(self):
+    def bs(self) -> int:
         """Returns the block size of the dofmap"""
         return self._cpp_object.bs
 
     @property
-    def dof_layout(self):
+    def dof_layout(self) -> ElementDofLayout:
         """Layout of dofs on an element."""
         return self._cpp_object.dof_layout
 
     @property
-    def index_map(self):
+    def index_map(self) -> IndexMap:
         """Index map that described the parallel distribution of the dofmap."""
         return self._cpp_object.index_map
 
     @property
-    def index_map_bs(self):
+    def index_map_bs(self) -> int:
         """Block size of the index map."""
         return self._cpp_object.index_map_bs
 
     @property
-    def list(self):
+    def list(self) -> npt.NDArray[np.int32]:
         """Adjacency list with dof indices for each cell."""
         return self._cpp_object.map()
diff --git a/python/dolfinx/fem/element.py b/python/dolfinx/fem/element.py
index 7e88427c12b..83df51f8d5c 100644
--- a/python/dolfinx/fem/element.py
+++ b/python/dolfinx/fem/element.py
@@ -51,9 +51,9 @@ def dtype(self) -> np.dtype:
 def coordinate_element(
     celltype: _cpp.mesh.CellType,
     degree: int,
-    variant=int(basix.LagrangeVariant.unset),
+    variant: int = 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.
@@ -76,7 +76,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.
diff --git a/python/dolfinx/fem/function.py b/python/dolfinx/fem/function.py
index a04ce990c67..fc499a93d5e 100644
--- a/python/dolfinx/fem/function.py
+++ b/python/dolfinx/fem/function.py
@@ -7,9 +7,10 @@
 
 from __future__ import annotations
 
-import typing
 import warnings
+from collections.abc import Sequence
 from functools import singledispatch
+from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -20,15 +21,16 @@
 from dolfinx import default_scalar_type, jit, la
 from dolfinx.fem import dofmap
 from dolfinx.geometry import PointOwnershipData
+from dolfinx.mesh import Mesh
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from mpi4py import MPI as _MPI
 
     from dolfinx.mesh import Mesh
 
 
 class Constant(ufl.Constant):
-    _cpp_object: typing.Union[
+    _cpp_object: Union[
         _cpp.fem.Constant_complex64,
         _cpp.fem.Constant_complex128,
         _cpp.fem.Constant_float32,
@@ -36,7 +38,9 @@ class Constant(ufl.Constant):
     ]
 
     def __init__(
-        self, domain, c: typing.Union[np.ndarray, typing.Sequence, np.floating, np.complexfloating]
+        self,
+        domain: Union[ufl.Mesh, Mesh],
+        c: Union[np.ndarray, Sequence, np.floating, np.complexfloating],
     ):
         """A constant with respect to a domain.
 
@@ -61,25 +65,25 @@ def __init__(
             raise AttributeError("Constant value must have a dtype attribute.")
 
     @property
-    def value(self):
+    def value(self) -> float:
         """The value of the constant"""
         return self._cpp_object.value
 
     @value.setter
-    def value(self, v):
+    def value(self, v: float) -> None:
         np.copyto(self._cpp_object.value, np.asarray(v))
 
     @property
     def dtype(self) -> np.dtype:
         return np.dtype(self._cpp_object.dtype)
 
-    def __float__(self):
+    def __float__(self) -> float:
         if self.ufl_shape or self.ufl_free_indices:
             raise TypeError("Cannot evaluate a nonscalar expression to a scalar value.")
         else:
             return float(self.value)
 
-    def __complex__(self):
+    def __complex__(self) -> complex:
         if self.ufl_shape or self.ufl_free_indices:
             raise TypeError("Cannot evaluate a nonscalar expression to a scalar value.")
         else:
@@ -91,10 +95,10 @@ def __init__(
         self,
         e: ufl.core.expr.Expr,
         X: np.ndarray,
-        comm: typing.Optional[_MPI.Comm] = None,
-        form_compiler_options: typing.Optional[dict] = None,
-        jit_options: typing.Optional[dict] = None,
-        dtype: typing.Optional[npt.DTypeLike] = None,
+        comm: Optional[_MPI.Comm] = None,
+        form_compiler_options: Optional[dict] = None,
+        jit_options: Optional[dict] = None,
+        dtype: Optional[npt.DTypeLike] = None,
     ):
         """Create a DOLFINx Expression.
 
@@ -173,7 +177,7 @@ def __init__(
         else:
             raise RuntimeError("Expressions with more that one Argument not allowed.")
 
-        def _create_expression(dtype):
+        def _create_expression(dtype: npt.DTypeLike) -> _cpp.fem.Expression:
             if np.issubdtype(dtype, np.float32):
                 return _cpp.fem.create_expression_float32
             elif np.issubdtype(dtype, np.float64):
@@ -197,7 +201,7 @@ def eval(
         self,
         mesh: Mesh,
         entities: np.ndarray,
-        values: typing.Optional[np.ndarray] = None,
+        values: Optional[np.ndarray] = None,
     ) -> np.ndarray:
         """Evaluate Expression on entities.
 
@@ -249,7 +253,7 @@ def X(self) -> np.ndarray:
         return self._cpp_object.X()
 
     @property
-    def ufl_expression(self):
+    def ufl_expression(self) -> Expression:
         """Original UFL Expression"""
         return self._ufl_expression
 
@@ -259,12 +263,12 @@ def value_size(self) -> int:
         return self._cpp_object.value_size
 
     @property
-    def argument_function_space(self) -> typing.Optional[FunctionSpace]:
+    def argument_function_space(self) -> Optional[FunctionSpace]:
         """The argument function space if expression has argument"""
         return self._argument_function_space
 
     @property
-    def ufcx_expression(self):
+    def ufcx_expression(self) -> tuple[Any, Any | None, tuple[None, None] | tuple[str, str] | Any]:
         """The compiled ufcx_expression object"""
         return self._ufcx_expression
 
@@ -283,7 +287,7 @@ class Function(ufl.Coefficient):
     (domain, element and dofmap) and a vector holding the
     degrees-of-freedom."""
 
-    _cpp_object: typing.Union[
+    _cpp_object: Union[
         _cpp.fem.Function_complex64,
         _cpp.fem.Function_complex128,
         _cpp.fem.Function_float32,
@@ -293,9 +297,9 @@ class Function(ufl.Coefficient):
     def __init__(
         self,
         V: FunctionSpace,
-        x: typing.Optional[la.Vector] = None,
-        name: typing.Optional[str] = None,
-        dtype: typing.Optional[npt.DTypeLike] = None,
+        x: Optional[la.Vector] = None,
+        name: Optional[str] = None,
+        dtype: Optional[npt.DTypeLike] = None,
     ):
         """Initialize a finite element Function.
 
@@ -321,7 +325,14 @@ def __init__(
         ), "Incompatible FunctionSpace dtype and requested dtype."
 
         # Create cpp Function
-        def functiontype(dtype):
+        def functiontype(
+            dtype: npt.DTypeLike,
+        ) -> Union[
+            _cpp.fem.Function_float32,
+            _cpp.fem.Function_float64,
+            _cpp.fem.Function_complex64,
+            _cpp.fem.Function_complex128,
+        ]:
             if np.issubdtype(dtype, np.float32):
                 return _cpp.fem.Function_float32
             elif np.issubdtype(dtype, np.float64):
@@ -358,7 +369,9 @@ def function_space(self) -> FunctionSpace:
         """The 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.ArrayLike, u: Optional[npt.NDArray[np.floating]] = None
+    ) -> npt.NDArray[np.floating]:
         """Evaluate Function at points x.
 
         Points where x has shape (num_points, 3), and cells has shape
@@ -411,9 +424,9 @@ def interpolate_nonmatching(
 
     def interpolate(
         self,
-        u0: typing.Union[typing.Callable, Expression, Function],
-        cells0: typing.Optional[np.ndarray] = None,
-        cells1: typing.Optional[np.ndarray] = None,
+        u0: Union[Callable, Expression, Function],
+        cells0: Optional[np.ndarray] = None,
+        cells1: Optional[np.ndarray] = None,
     ) -> None:
         """Interpolate an expression.
 
@@ -436,22 +449,22 @@ def interpolate(
             cells1 = np.arange(0, dtype=np.int32)
 
         @singledispatch
-        def _interpolate(u0):
+        def _interpolate(u0: Any) -> None:
             """Interpolate a cpp.fem.Function."""
             self._cpp_object.interpolate(u0, cells0, cells1)  # type: ignore
 
         @_interpolate.register(Function)
-        def _(u0: Function):
+        def _(u0: Function) -> None:
             """Interpolate a fem.Function."""
             self._cpp_object.interpolate(u0._cpp_object, cells0, cells1)  # type: ignore
 
         @_interpolate.register(int)
-        def _(u0_ptr: int):
+        def _(u0_ptr: int) -> None:
             """Interpolate using a pointer to a function f(x)."""
             self._cpp_object.interpolate_ptr(u0_ptr, cells0)  # type: ignore
 
         @_interpolate.register(Expression)
-        def _(e0: Expression):
+        def _(e0: Expression) -> None:
             """Interpolate a fem.Expression."""
             self._cpp_object.interpolate(e0._cpp_object, cells0, cells1)  # type: ignore
 
@@ -483,7 +496,7 @@ def x(self) -> la.Vector:
         return self._x
 
     @property
-    def vector(self):
+    def vector(self) -> Any:
         """PETSc vector holding the degrees-of-freedom.
 
         Upon first call, this function creates a PETSc ``Vec`` object
@@ -511,10 +524,10 @@ def name(self) -> str:
         return self._cpp_object.name  # type: ignore
 
     @name.setter
-    def name(self, name):
+    def name(self, name: str) -> None:
         self._cpp_object.name = name
 
-    def __str__(self):
+    def __str__(self) -> str:
         """Pretty print representation."""
         return self.name
 
@@ -562,7 +575,7 @@ def collapse(self) -> Function:
         return Function(V_collapsed, la.Vector(u_collapsed.x))
 
 
-class ElementMetaData(typing.NamedTuple):
+class ElementMetaData(NamedTuple):
     """Data for representing a finite element
 
     :param family: Element type.
@@ -574,16 +587,16 @@ class ElementMetaData(typing.NamedTuple):
 
     family: str
     degree: int
-    shape: typing.Optional[tuple[int, ...]] = None
-    symmetry: typing.Optional[bool] = None
+    shape: Optional[tuple[int, ...]] = None
+    symmetry: Optional[bool] = None
 
 
 def _create_dolfinx_element(
-    comm: _MPI.Intracomm,
+    comm: _MPI.Comm,
     cell_type: _cpp.mesh.CellType,
     ufl_e: ufl.FiniteElementBase,
     dtype: np.dtype,
-) -> typing.Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]:
+) -> Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]:
     """Create a DOLFINx element from a basix.ufl element."""
     if np.issubdtype(dtype, np.float32):
         CppElement = _cpp.fem.FiniteElement_float32
@@ -606,9 +619,9 @@ def _create_dolfinx_element(
 
 def functionspace(
     mesh: Mesh,
-    element: typing.Union[ufl.FiniteElementBase, ElementMetaData, tuple[str, int, tuple, bool]],
-    form_compiler_options: typing.Optional[dict[str, typing.Any]] = None,
-    jit_options: typing.Optional[dict[str, typing.Any]] = None,
+    element: Union[ufl.FiniteElementBase, ElementMetaData, tuple[str, int, tuple, bool]],
+    form_compiler_options: Optional[dict[str, Any]] = None,
+    jit_options: Optional[dict[str, Any]] = None,
 ) -> FunctionSpace:
     """Create a finite element function space.
 
@@ -667,14 +680,14 @@ def functionspace(
 class FunctionSpace(ufl.FunctionSpace):
     """A space on which Functions (fields) can be defined."""
 
-    _cpp_object: typing.Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64]
+    _cpp_object: Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64]
     _mesh: Mesh
 
     def __init__(
         self,
         mesh: Mesh,
         element: ufl.FiniteElementBase,
-        cppV: typing.Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64],
+        cppV: Union[_cpp.fem.FunctionSpace_float32, _cpp.fem.FunctionSpace_float64],
     ):
         """Create a finite element function space.
 
@@ -756,11 +769,11 @@ def sub(self, i: int) -> FunctionSpace:
         cppV_sub = self._cpp_object.sub([i])  # type: ignore
         return FunctionSpace(self._mesh, sub_element, cppV_sub)
 
-    def component(self):
+    def component(self) -> npt.NDArray[np.integer]:
         """Return the component relative to the parent space."""
-        return self._cpp_object.component()  # type: ignore
+        return self._cpp_object.component()
 
-    def contains(self, V) -> bool:
+    def contains(self, V: FunctionSpace) -> bool:
         """Check if a space is contained in, or is the same as
         (identity), this space.
 
@@ -773,11 +786,11 @@ def contains(self, V) -> bool:
         """
         return self._cpp_object.contains(V._cpp_object)  # type: ignore
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         """Comparison for equality."""
         return super().__eq__(other) and self._cpp_object == other._cpp_object
 
-    def __ne__(self, other):
+    def __ne__(self, other: Any) -> bool:
         """Comparison for inequality."""
         return super().__ne__(other) or self._cpp_object != other._cpp_object
 
@@ -788,7 +801,7 @@ def ufl_function_space(self) -> ufl.FunctionSpace:
     @property
     def element(
         self,
-    ) -> typing.Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]:
+    ) -> Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]:
         """Function space finite element."""
         return self._cpp_object.element  # type: ignore
 
diff --git a/python/dolfinx/geometry.py b/python/dolfinx/geometry.py
index 03cd0870b24..84794520611 100644
--- a/python/dolfinx/geometry.py
+++ b/python/dolfinx/geometry.py
@@ -9,11 +9,13 @@
 
 import typing
 
+from mpi4py import MPI as _MPI
+
 import numpy as np
 import numpy.typing as npt
 
 if typing.TYPE_CHECKING:
-    from dolfinx.cpp.graph import AdjacencyList_int32
+    from dolfinx.cpp.graph import AdjacencyList_int32, AdjacencyList_int64
     from dolfinx.mesh import Mesh
 
 
@@ -40,7 +42,12 @@ class PointOwnershipData:
         _cpp.geometry.PointOwnershipData_float32, _cpp.geometry.PointOwnershipData_float64
     ]
 
-    def __init__(self, ownership_data):
+    def __init__(
+        self,
+        ownership_data: typing.Union[
+            _cpp.geometry.PointOwnershipData_float32, _cpp.geometry.PointOwnershipData_float64
+        ],
+    ) -> None:
         """Wrap a C++ PointOwnershipData."""
         self._cpp_object = ownership_data
 
@@ -68,7 +75,12 @@ class BoundingBoxTree:
         _cpp.geometry.BoundingBoxTree_float32, _cpp.geometry.BoundingBoxTree_float64
     ]
 
-    def __init__(self, tree):
+    def __init__(
+        self,
+        tree: typing.Union[
+            _cpp.geometry.BoundingBoxTree_float32, _cpp.geometry.BoundingBoxTree_float64
+        ],
+    ) -> None:
         """Wrap a C++ BoundingBoxTree.
 
         Note:
@@ -83,7 +95,7 @@ def num_bboxes(self) -> int:
         """Number of bounding boxes."""
         return self._cpp_object.num_bboxes
 
-    def get_bbox(self, i) -> npt.NDArray[np.floating]:
+    def get_bbox(self, i: int) -> npt.NDArray[np.floating]:
         """Get lower and upper corners of the ith bounding box.
 
         Args:
@@ -96,7 +108,7 @@ def get_bbox(self, i) -> npt.NDArray[np.floating]:
         """
         return self._cpp_object.get_bbox(i)
 
-    def create_global_tree(self, comm) -> BoundingBoxTree:
+    def create_global_tree(self, comm: _MPI.Comm) -> BoundingBoxTree:
         return BoundingBoxTree(self._cpp_object.create_global_tree(comm))
 
 
@@ -217,7 +229,7 @@ def create_midpoint_tree(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32])
 
 def compute_colliding_cells(
     mesh: Mesh, candidates: AdjacencyList_int32, x: npt.NDArray[np.floating]
-):
+) -> typing.Union[AdjacencyList_int32, AdjacencyList_int64]:
     """From a mesh, find which cells collide with a set of points.
 
     Args:
@@ -234,7 +246,9 @@ def compute_colliding_cells(
     return _cpp.geometry.compute_colliding_cells(mesh._cpp_object, candidates, x)
 
 
-def squared_distance(mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating]):
+def squared_distance(
+    mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """Compute the squared distance between a point and a mesh entity.
 
     The distance is computed between the ith input points and the ith
diff --git a/python/dolfinx/graph.py b/python/dolfinx/graph.py
index dd16514a79a..18714b71540 100644
--- a/python/dolfinx/graph.py
+++ b/python/dolfinx/graph.py
@@ -7,7 +7,10 @@
 
 from __future__ import annotations
 
+from typing import Optional, Union
+
 import numpy as np
+import numpy.typing as npt
 
 from dolfinx import cpp as _cpp
 from dolfinx.cpp.graph import partitioner
@@ -31,7 +34,9 @@
 __all__ = ["adjacencylist", "partitioner"]
 
 
-def adjacencylist(data: np.ndarray, offsets=None):
+def adjacencylist(
+    data: np.ndarray, offsets: Optional[npt.NDArray[Union[np.int32, np.int64]]] = None
+) -> Union[_cpp.graph.AdjacencyList_int32, _cpp.graph.AdjacencyList_int64]:
     """Create an AdjacencyList for int32 or int64 datasets.
 
     Args:
diff --git a/python/dolfinx/io/utils.py b/python/dolfinx/io/utils.py
index a479f898998..e1ae5dcb159 100644
--- a/python/dolfinx/io/utils.py
+++ b/python/dolfinx/io/utils.py
@@ -8,6 +8,7 @@
 
 import typing
 from pathlib import Path
+from types import TracebackType
 
 from mpi4py import MPI as _MPI
 
@@ -26,11 +27,13 @@
 __all__ = ["VTKFile", "XDMFFile", "cell_perm_gmsh", "cell_perm_vtk", "distribute_entity_data"]
 
 
-def _extract_cpp_objects(functions: typing.Union[Mesh, Function, tuple[Function], list[Function]]):
+def _extract_cpp_objects(
+    functions: typing.Union[Mesh, Function, tuple[Function], list[Function]],
+) -> list[typing.Any | Function]:
     """Extract C++ objects"""
     if isinstance(functions, (list, tuple)):
         return [getattr(u, "_cpp_object", u) for u in functions]
-    else:
+    else:  # TODO: remove case can not happen as indicated by typing
         return [getattr(functions, "_cpp_object", functions)]
 
 
@@ -103,16 +106,21 @@ def __init__(
                     comm, filename, _extract_cpp_objects(output), engine, mesh_policy
                 )  # type: ignore[arg-type]
 
-        def __enter__(self):
+        def __enter__(self) -> typing.Self:
             return self
 
-        def __exit__(self, exception_type, exception_value, traceback):
+        def __exit__(
+            self,
+            exception_type: typing.Optional[BaseException],
+            exception_value: typing.Optional[BaseException],
+            traceback: typing.Optional[TracebackType],
+        ) -> None:
             self.close()
 
-        def write(self, t: float):
+        def write(self, t: float) -> None:
             self._cpp_object.write(t)
 
-        def close(self):
+        def close(self) -> None:
             self._cpp_object.close()
 
     class FidesWriter:
@@ -174,16 +182,21 @@ def __init__(
                     comm, filename, _extract_cpp_objects(output), engine, mesh_policy
                 )  # type: ignore[arg-type]
 
-        def __enter__(self):
+        def __enter__(self) -> typing.Self:
             return self
 
-        def __exit__(self, exception_type, exception_value, traceback):
+        def __exit__(
+            self,
+            exception_type: typing.Optional[BaseException],
+            exception_value: typing.Optional[BaseException],
+            traceback: typing.Optional[TracebackType],
+        ) -> None:
             self.close()
 
-        def write(self, t: float):
+        def write(self, t: float) -> None:
             self._cpp_object.write(t)
 
-        def close(self):
+        def close(self) -> None:
             self._cpp_object.close()
 
 
@@ -196,10 +209,15 @@ class VTKFile(_cpp.io.VTKFile):
 
     """
 
-    def __enter__(self):
+    def __enter__(self) -> typing.Self:
         return self
 
-    def __exit__(self, exception_type, exception_value, traceback):
+    def __exit__(
+        self,
+        exception_type: typing.Optional[BaseException],
+        exception_value: typing.Optional[BaseException],
+        traceback: typing.Optional[TracebackType],
+    ) -> None:
         self.close()
 
     def write_mesh(self, mesh: Mesh, t: float = 0.0) -> None:
@@ -212,10 +230,15 @@ def write_function(self, u: typing.Union[list[Function], Function], t: float = 0
 
 
 class XDMFFile(_cpp.io.XDMFFile):
-    def __enter__(self):
+    def __enter__(self) -> typing.Self:
         return self
 
-    def __exit__(self, exception_type, exception_value, traceback):
+    def __exit__(
+        self,
+        exception_type: typing.Optional[BaseException],
+        exception_value: typing.Optional[BaseException],
+        traceback: typing.Optional[TracebackType],
+    ) -> None:
         self.close()
 
     def write_mesh(self, mesh: Mesh, xpath: str = "/Xdmf/Domain") -> None:
@@ -233,8 +256,11 @@ def write_meshtags(
         super().write_meshtags(tags._cpp_object, x, geometry_xpath, xpath)
 
     def write_function(
-        self, u: Function, t: float = 0.0, mesh_xpath="/Xdmf/Domain/Grid[@GridType='Uniform'][1]"
-    ):
+        self,
+        u: Function,
+        t: float = 0.0,
+        mesh_xpath: str = "/Xdmf/Domain/Grid[@GridType='Uniform'][1]",
+    ) -> None:
         """Write function to file for a given time.
 
         Note:
@@ -252,7 +278,10 @@ def write_function(
         super().write_function(getattr(u, "_cpp_object", u), t, mesh_xpath)
 
     def read_mesh(
-        self, ghost_mode=GhostMode.shared_facet, name="mesh", xpath="/Xdmf/Domain"
+        self,
+        ghost_mode: GhostMode = GhostMode.shared_facet,
+        name: str = "mesh",
+        xpath: str = "/Xdmf/Domain",
     ) -> Mesh:
         """Read mesh data from file."""
         cell_shape, cell_degree = super().read_cell_type(name, xpath)
@@ -276,7 +305,7 @@ def read_mesh(
         )
         return Mesh(msh, domain)
 
-    def read_meshtags(self, mesh, name, xpath="/Xdmf/Domain"):
+    def read_meshtags(self, mesh: Mesh, name: str, xpath: str = "/Xdmf/Domain") -> MeshTags:
         mt = super().read_meshtags(mesh._cpp_object, name, xpath)
         return MeshTags(mt)
 
diff --git a/python/dolfinx/jit.py b/python/dolfinx/jit.py
index d27093599bc..9bb37c32fda 100644
--- a/python/dolfinx/jit.py
+++ b/python/dolfinx/jit.py
@@ -10,9 +10,9 @@
 import os
 import sys
 from pathlib import Path
-from typing import Optional
+from typing import Any, Callable, Optional
 
-from mpi4py import MPI
+from mpi4py import MPI as _MPI
 
 import ffcx
 import ffcx.codegeneration.jit
@@ -44,7 +44,9 @@
     )
 
 
-def mpi_jit_decorator(local_jit, *args, **kwargs):
+def mpi_jit_decorator(
+    local_jit: Callable[..., Any], *args: Any, **kwargs: Any
+) -> Callable[..., Any]:
     """A decorator for jit compilation.
 
     Use this function as a decorator to any jit compiler function. In a
@@ -56,7 +58,7 @@ def mpi_jit_decorator(local_jit, *args, **kwargs):
     """
 
     @functools.wraps(local_jit)
-    def mpi_jit(comm, *args, **kwargs):
+    def mpi_jit(comm: _MPI.Comm, *args: Any, **kwargs: Any) -> Callable[..., Any]:
         # Just call JIT compiler when running in serial
         if comm.size == 1:
             return local_jit(*args, **kwargs)
@@ -83,7 +85,7 @@ def mpi_jit(comm, *args, **kwargs):
         # Wait for the compiling process to finish and get status
         # TODO: Would be better to broadcast the status from root but
         # this works.
-        global_status = comm.allreduce(status, op=MPI.MAX)
+        global_status = comm.allreduce(status, op=_MPI.MAX)
         if global_status == 0:
             # Success, call jit on all other processes (this should just
             # read the cache)
@@ -103,7 +105,7 @@ def mpi_jit(comm, *args, **kwargs):
 
 
 @functools.cache
-def _load_options():
+def _load_options() -> tuple[dict, dict]:
     """Loads options from JSON files."""
     user_config_file = os.getenv("XDG_CONFIG_HOME", default=Path.home().joinpath(".config")) / Path(
         "dolfinx", "dolfinx_jit_options.json"
@@ -157,8 +159,10 @@ def get_options(priority_options: Optional[dict] = None) -> dict:
 
 @mpi_jit_decorator
 def ffcx_jit(
-    ufl_object, form_compiler_options: Optional[dict] = None, jit_options: Optional[dict] = None
-):
+    ufl_object: ufl.Form,
+    form_compiler_options: Optional[dict] = None,
+    jit_options: Optional[dict] = None,
+) -> tuple[Any, Any, tuple[str, str]]:
     """Compile UFL object with FFCx and CFFI.
 
     Args:
diff --git a/python/dolfinx/la.py b/python/dolfinx/la.py
index 1c7b58ab676..cf76905e112 100644
--- a/python/dolfinx/la.py
+++ b/python/dolfinx/la.py
@@ -5,7 +5,7 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 """Linear algebra functionality"""
 
-import typing
+from typing import Any, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -28,7 +28,7 @@
 
 
 class MatrixCSR:
-    _cpp_object: typing.Union[
+    _cpp_object: Union[
         _cpp.la.MatrixCSR_float32,
         _cpp.la.MatrixCSR_float64,
         _cpp.la.MatrixCSR_complex64,
@@ -37,7 +37,7 @@ class MatrixCSR:
 
     def __init__(
         self,
-        A: typing.Union[
+        A: Union[
             _cpp.la.MatrixCSR_float32,
             _cpp.la.MatrixCSR_float64,
             _cpp.la.MatrixCSR_complex64,
@@ -64,7 +64,7 @@ def index_map(self, i: int) -> IndexMap:
         return self._cpp_object.index_map(i)
 
     @property
-    def block_size(self):
+    def block_size(self) -> tuple[int, int]:
         """Block sizes for the matrix."""
         return self._cpp_object.bs
 
@@ -132,7 +132,7 @@ def to_dense(self) -> npt.NDArray[np.floating]:
         """
         return self._cpp_object.to_dense()
 
-    def to_scipy(self, ghosted=False):
+    def to_scipy(self, ghosted: bool = False) -> Any:
         """Convert to a SciPy CSR/BSR matrix. Data is shared.
 
         Note:
@@ -174,7 +174,9 @@ def to_scipy(self, ghosted=False):
 
 
 def matrix_csr(
-    sp: _cpp.la.SparsityPattern, block_mode=BlockMode.compact, dtype: npt.DTypeLike = np.float64
+    sp: _cpp.la.SparsityPattern,
+    block_mode: BlockMode = BlockMode.compact,
+    dtype: npt.DTypeLike = np.float64,
 ) -> MatrixCSR:
     """Create a distributed sparse matrix.
 
@@ -203,7 +205,7 @@ def matrix_csr(
 
 
 class Vector:
-    _cpp_object: typing.Union[
+    _cpp_object: Union[
         _cpp.la.Vector_float32,
         _cpp.la.Vector_float64,
         _cpp.la.Vector_complex64,
@@ -215,7 +217,7 @@ class Vector:
 
     def __init__(
         self,
-        x: typing.Union[
+        x: Union[
             _cpp.la.Vector_float32,
             _cpp.la.Vector_float64,
             _cpp.la.Vector_complex64,
@@ -237,7 +239,7 @@ def __init__(
         self._cpp_object = x
         self._petsc_x = None
 
-    def __del__(self):
+    def __del__(self) -> None:
         if self._petsc_x is not None:
             self._petsc_x.destroy()
 
@@ -257,7 +259,7 @@ def array(self) -> np.ndarray:
         return self._cpp_object.array
 
     @property
-    def petsc_vec(self):
+    def petsc_vec(self) -> Any:
         """PETSc vector holding the entries of the vector.
 
         Upon first call, this function creates a PETSc ``Vec`` object
@@ -282,7 +284,7 @@ def scatter_reverse(self, mode: InsertMode) -> None:
         self._cpp_object.scatter_reverse(mode)
 
 
-def vector(map, bs=1, dtype: npt.DTypeLike = np.float64) -> Vector:
+def vector(map: IndexMap, bs: int = 1, dtype: npt.DTypeLike = np.float64) -> Vector:
     """Create a distributed vector.
 
     Args:
@@ -314,7 +316,7 @@ def vector(map, bs=1, dtype: npt.DTypeLike = np.float64) -> Vector:
     return Vector(vtype(map, bs))
 
 
-def create_petsc_vector_wrap(x: Vector):
+def create_petsc_vector_wrap(x: Vector) -> Any:
     """Wrap a distributed DOLFINx vector as a PETSc vector.
 
     Args:
@@ -336,7 +338,7 @@ def create_petsc_vector_wrap(x: Vector):
     return PETSc.Vec().createGhostWithArray(ghosts, x.array, size=size, bsize=bs, comm=map.comm)  # type: ignore
 
 
-def create_petsc_vector(map, bs: int):
+def create_petsc_vector(map: IndexMap, bs: int) -> Any:
     """Create a distributed PETSc vector.
 
     Args:
@@ -354,7 +356,7 @@ def create_petsc_vector(map, bs: int):
     return PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=map.comm)  # type: ignore
 
 
-def orthonormalize(basis: list[Vector]):
+def orthonormalize(basis: list[Vector]) -> None:
     """Orthogonalise set of PETSc vectors in-place."""
     _cpp.la.orthonormalize([x._cpp_object for x in basis])
 
diff --git a/python/dolfinx/mesh.py b/python/dolfinx/mesh.py
index 5f1a0da94bb..7a800f4a410 100644
--- a/python/dolfinx/mesh.py
+++ b/python/dolfinx/mesh.py
@@ -19,10 +19,18 @@
 import ufl
 from dolfinx import cpp as _cpp
 from dolfinx import default_real_type
+from dolfinx.cpp.graph import AdjacencyList_int32, AdjacencyList_int64
 from dolfinx.cpp.mesh import (
     CellType,
     DiagonalType,
+    Geometry_float32,
+    Geometry_float64,
     GhostMode,
+    MeshTags_float64,
+    MeshTags_int8,
+    MeshTags_int32,
+    MeshTags_int64,
+    Topology,
     build_dual_graph,
     cell_dim,
     create_cell_partitioner,
@@ -66,6 +74,11 @@
     "entities_to_geometry",
 ]
 
+CellPartionerType = typing.Callable[
+    [_MPI.Comm, int, list[CellType], list[npt.NDArray[np.int64]]],
+    AdjacencyList_int32 | AdjacencyList_int64,
+]
+
 
 class Mesh:
     """A mesh."""
@@ -73,7 +86,11 @@ class Mesh:
     _mesh: typing.Union[_cpp.mesh.Mesh_float32, _cpp.mesh.Mesh_float64]
     _ufl_domain: typing.Optional[ufl.Mesh]
 
-    def __init__(self, mesh, domain: typing.Optional[ufl.Mesh]):
+    def __init__(
+        self,
+        mesh: typing.Union[_cpp.mesh.Mesh_float32, _cpp.mesh.Mesh_float64],
+        domain: typing.Optional[ufl.Mesh],
+    ):
         """Initialize mesh from a C++ mesh.
 
         Args:
@@ -90,15 +107,15 @@ def __init__(self, mesh, domain: typing.Optional[ufl.Mesh]):
             self._ufl_domain._ufl_cargo = self._cpp_object  # type: ignore
 
     @property
-    def comm(self):
+    def comm(self) -> _MPI.Comm:
         return self._cpp_object.comm
 
     @property
-    def name(self):
+    def name(self) -> str:
         return self._cpp_object.name
 
     @name.setter
-    def name(self, value):
+    def name(self, value: str) -> None:
         self._cpp_object.name = value
 
     def ufl_cell(self) -> ufl.Cell:
@@ -138,12 +155,12 @@ def h(self, dim: int, entities: npt.NDArray[np.int32]) -> npt.NDArray[np.float64
         return _cpp.mesh.h(self._cpp_object, dim, entities)
 
     @property
-    def topology(self):
+    def topology(self) -> Topology:
         "Mesh topology."
         return self._cpp_object.topology
 
     @property
-    def geometry(self):
+    def geometry(self) -> typing.Union[Geometry_float32, Geometry_float64]:
         "Mesh geometry."
         return self._cpp_object.geometry
 
@@ -151,7 +168,10 @@ def geometry(self):
 class MeshTags:
     """Mesh tags associate data (markers) with a subset of mesh entities of a given dimension."""
 
-    def __init__(self, meshtags):
+    def __init__(
+        self,
+        meshtags: typing.Union[MeshTags_int8, MeshTags_int32, MeshTags_int64, MeshTags_float64],
+    ):
         """Initialize tags from a C++ MeshTags object.
 
         Args:
@@ -189,7 +209,7 @@ def indices(self) -> npt.NDArray[np.int32]:
         return self._cpp_object.indices
 
     @property
-    def values(self):
+    def values(self) -> npt.NDArray[np.int32]:
         """Values associated with tagged mesh entities."""
         return self._cpp_object.values
 
@@ -199,10 +219,10 @@ def name(self) -> str:
         return self._cpp_object.name
 
     @name.setter
-    def name(self, value):
+    def name(self, value: np.int32) -> None:
         self._cpp_object.name = value
 
-    def find(self, value) -> npt.NDArray[np.int32]:
+    def find(self, value: np.int32) -> npt.NDArray[np.int32]:
         """Get a list of all entity indices with a given value.
 
         Args:
@@ -214,11 +234,15 @@ def find(self, value) -> npt.NDArray[np.int32]:
         return self._cpp_object.find(value)
 
 
-def compute_incident_entities(topology, entities: npt.NDArray[np.int32], d0: int, d1: int):
+def compute_incident_entities(
+    topology: Topology, entities: npt.NDArray[np.int32], d0: int, d1: int
+) -> npt.NDArray[np.int32]:
     return _cpp.mesh.compute_incident_entities(topology, entities, d0, d1)
 
 
-def compute_midpoints(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32]):
+def compute_midpoints(
+    mesh: Mesh, dim: int, entities: npt.NDArray[np.int32]
+) -> npt.NDArray[np.floating]:
     return _cpp.mesh.compute_midpoints(mesh._cpp_object, dim, entities)
 
 
@@ -550,7 +574,7 @@ def meshtags(
 
 def meshtags_from_entities(
     mesh: Mesh, dim: int, entities: _cpp.graph.AdjacencyList_int32, values: npt.NDArray[typing.Any]
-):
+) -> MeshTags:
     """Create a :class:dolfinx.mesh.MeshTags` object that associates
     data with a subset of mesh entities, where the entities are defined
     by their vertices.
@@ -584,8 +608,8 @@ def create_interval(
     nx: int,
     points: npt.ArrayLike,
     dtype: npt.DTypeLike = default_real_type,
-    ghost_mode=GhostMode.shared_facet,
-    partitioner=None,
+    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[CellPartionerType] = None,
 ) -> Mesh:
     """Create an interval mesh.
 
@@ -620,8 +644,8 @@ def create_unit_interval(
     comm: _MPI.Comm,
     nx: int,
     dtype: npt.DTypeLike = default_real_type,
-    ghost_mode=GhostMode.shared_facet,
-    partitioner=None,
+    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[CellPartionerType] = None,
 ) -> Mesh:
     """Create a mesh on the unit interval.
 
@@ -646,10 +670,10 @@ def create_rectangle(
     comm: _MPI.Comm,
     points: npt.ArrayLike,
     n: npt.ArrayLike,
-    cell_type=CellType.triangle,
+    cell_type: CellType = CellType.triangle,
     dtype: npt.DTypeLike = default_real_type,
-    ghost_mode=GhostMode.shared_facet,
-    partitioner=None,
+    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[CellPartionerType] = None,
     diagonal: DiagonalType = DiagonalType.right,
 ) -> Mesh:
     """Create a rectangle mesh.
@@ -688,10 +712,10 @@ def create_unit_square(
     comm: _MPI.Comm,
     nx: int,
     ny: int,
-    cell_type=CellType.triangle,
+    cell_type: CellType = CellType.triangle,
     dtype: npt.DTypeLike = default_real_type,
-    ghost_mode=GhostMode.shared_facet,
-    partitioner=None,
+    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[CellPartionerType] = None,
     diagonal: DiagonalType = DiagonalType.right,
 ) -> Mesh:
     """Create a mesh of a unit square.
@@ -728,10 +752,10 @@ def create_box(
     comm: _MPI.Comm,
     points: list[npt.ArrayLike],
     n: list,
-    cell_type=CellType.tetrahedron,
+    cell_type: CellType = CellType.tetrahedron,
     dtype: npt.DTypeLike = default_real_type,
-    ghost_mode=GhostMode.shared_facet,
-    partitioner=None,
+    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[CellPartionerType] = None,
 ) -> Mesh:
     """Create a box mesh.
 
@@ -767,10 +791,10 @@ def create_unit_cube(
     nx: int,
     ny: int,
     nz: int,
-    cell_type=CellType.tetrahedron,
+    cell_type: CellType = CellType.tetrahedron,
     dtype: npt.DTypeLike = default_real_type,
-    ghost_mode=GhostMode.shared_facet,
-    partitioner=None,
+    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[CellPartionerType] = None,
 ) -> Mesh:
     """Create a mesh of a unit cube.
 
@@ -802,7 +826,7 @@ def create_unit_cube(
 
 
 def entities_to_geometry(
-    mesh: Mesh, dim: int, entities: npt.NDArray[np.int32], permute=False
+    mesh: Mesh, dim: int, entities: npt.NDArray[np.int32], permute: bool = False
 ) -> npt.NDArray[np.int32]:
     """Compute the geometric DOFs associated with the closure of the given mesh entities.
 
diff --git a/python/dolfinx/nls/petsc.py b/python/dolfinx/nls/petsc.py
index 6400034d83d..888c21c780d 100644
--- a/python/dolfinx/nls/petsc.py
+++ b/python/dolfinx/nls/petsc.py
@@ -37,11 +37,11 @@ def __init__(self, comm: MPI.Intracomm, problem: NonlinearProblem):
         self.setF(problem.F, self._b)
         self.set_form(problem.form)
 
-    def __del__(self):
+    def __del__(self) -> None:
         self._A.destroy()
         self._b.destroy()
 
-    def solve(self, u: fem.Function):
+    def solve(self, u: fem.Function) -> tuple[int, bool]:
         """Solve non-linear problem into function u. Returns the number
         of iterations and if the solver converged."""
         n, converged = super().solve(u.x.petsc_vec)
diff --git a/python/dolfinx/pkgconfig.py b/python/dolfinx/pkgconfig.py
index 6d4f9e0fe4e..075f3941a57 100644
--- a/python/dolfinx/pkgconfig.py
+++ b/python/dolfinx/pkgconfig.py
@@ -14,7 +14,7 @@
 import subprocess
 
 
-def _pkgconfig_query(s):
+def _pkgconfig_query(s: str) -> tuple[int, str]:
     pkg_config_exe = os.environ.get("PKG_CONFIG", None) or "pkg-config"
     cmd = [pkg_config_exe, *s.split()]
     proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
@@ -23,12 +23,12 @@ def _pkgconfig_query(s):
     return (rc, out.rstrip().decode("utf-8"))
 
 
-def exists(package) -> bool:
+def exists(package: str) -> bool:
     """Test for the existence of a pkg-config file for a named package."""
     return _pkgconfig_query("--exists " + package)[0] == 0
 
 
-def parse(package):
+def parse(package: str) -> dict[str, list[str]]:
     """Return a dict containing compile-time definitions."""
     parse_map = {
         "-D": "define_macros",
diff --git a/python/dolfinx/plot.py b/python/dolfinx/plot.py
index ef90cb1c2ca..c57395e4d14 100644
--- a/python/dolfinx/plot.py
+++ b/python/dolfinx/plot.py
@@ -6,9 +6,10 @@
 """Support functions for plotting"""
 
 import functools
-import typing
+from typing import Optional
 
 import numpy as np
+import numpy.typing as npt
 
 from dolfinx import cpp as _cpp
 from dolfinx import fem, mesh
@@ -31,7 +32,9 @@
 
 
 @functools.singledispatch
-def vtk_mesh(msh: mesh.Mesh, dim: typing.Optional[int] = None, entities=None):
+def vtk_mesh(
+    msh: mesh.Mesh, dim: Optional[int] = None, entities: Optional[npt.NDArray[np.int32]] = None
+) -> tuple[mesh.Topology, npt.NDArray[np.int8], npt.NDArray[np.floating]]:
     """Create vtk mesh topology data for mesh entities of a given
     dimension. The vertex indices in the returned topology array are the
     indices for the associated entry in the mesh geometry.
@@ -75,7 +78,9 @@ def vtk_mesh(msh: mesh.Mesh, dim: typing.Optional[int] = None, entities=None):
 
 
 @vtk_mesh.register(fem.FunctionSpace)
-def _(V: fem.FunctionSpace, entities=None):
+def _(
+    V: fem.FunctionSpace, entities: Optional[npt.NDArray[np.int8] | range] = None
+) -> tuple[mesh.Topology, npt.NDArray[np.int8], npt.NDArray[np.floating]]:
     """Creates a VTK mesh topology (topology array and array of cell
     types) that is based on the degree-of-freedom coordinates.
 
diff --git a/python/dolfinx/wrappers/__init__.py b/python/dolfinx/wrappers/__init__.py
index 9757a1d6da5..54aadac0876 100644
--- a/python/dolfinx/wrappers/__init__.py
+++ b/python/dolfinx/wrappers/__init__.py
@@ -4,9 +4,9 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
+from pathlib import Path
 
-def get_include_path():
-    """Return path to nanobind wrapper header files"""
-    import pathlib
 
-    return pathlib.Path(__file__).parent
+def get_include_path() -> Path:
+    """Return path to nanobind wrapper header files"""
+    return Path(__file__).parent
diff --git a/python/pyproject.toml b/python/pyproject.toml
index ef5a514e36f..ad58679e5dc 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -71,7 +71,7 @@ markers = [
 [tool.mypy]
 # Suggested at https://blog.wolt.com/engineering/2021/09/30/professional-grade-mypy-configuration/
 # Goal would be to make all of the below True long-term
-disallow_untyped_defs = false
+disallow_untyped_defs = true
 disallow_any_unimported = false
 no_implicit_optional = true
 check_untyped_defs = false