From 641fc5621c097115d274a740dc74a541a2477b9f Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Mon, 31 Aug 2026 15:49:34 +0200
Subject: [PATCH] Use MPI.Comm as type annotation in main functions

---
 docs/adding_backend.md          |  4 ++--
 src/io4dolfinx/checkpointing.py | 10 +++++-----
 tests/conftest.py               |  2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/docs/adding_backend.md b/docs/adding_backend.md
index 8017783..350ec35 100644
--- a/docs/adding_backend.md
+++ b/docs/adding_backend.md
@@ -21,7 +21,7 @@ You should import these to type-hint your implementation correctly:
 
 ## Implementation Checklist
 
-Your backend module must implement the functions listed below. Note that `comm` is always an {py:class}`MPI.Intracomm<mpi4py.MPI.Intracomm>` and `filename` is a
+Your backend module must implement the functions listed below. Note that `comm` is always an {py:class}`MPI.Comm<mpi4py.MPI.Comm>` and `filename` is a
 {py:class}`pathlib.Path` or {py:class}`str`.
 
 ### General Configuration
@@ -96,7 +96,7 @@ def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]
 
 def write_mesh(
     filename: Path | str,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     mesh: MeshData,
     backend_args: dict[str, Any] | None,
     mode: FileMode,
diff --git a/src/io4dolfinx/checkpointing.py b/src/io4dolfinx/checkpointing.py
index a128ef6..21f42c4 100644
--- a/src/io4dolfinx/checkpointing.py
+++ b/src/io4dolfinx/checkpointing.py
@@ -56,7 +56,7 @@
 
 def write_attributes(
     filename: Path | str,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     name: str,
     attributes: dict[str, np.ndarray],
     backend_args: dict[str, typing.Any] | None = None,
@@ -81,7 +81,7 @@ def write_attributes(
 
 def read_attributes(
     filename: Path | str,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     name: str,
     backend_args: dict[str, typing.Any] | None = None,
     backend: str | None = None,
@@ -106,7 +106,7 @@ def read_attributes(
 
 def read_timestamps(
     filename: Path | str,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     function_name: str,
     backend_args: dict[str, typing.Any] | None = None,
     backend: str | None = None,
@@ -402,7 +402,7 @@ def read_function(
 
 def read_mesh(
     filename: Path | str,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet,
     time: float | str | None = 0.0,
     read_from_partition: bool = False,
@@ -620,7 +620,7 @@ def write_function(
 
 def read_function_names(
     filename: Path | str,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     backend_args: dict[str, Any] | None = None,
     backend: str = "h5py",
 ) -> list[str]:
diff --git a/tests/conftest.py b/tests/conftest.py
index c3d6a82..3e6a475 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -216,7 +216,7 @@ def _read_function_time_dep(
 def _generate_reference_map(
     mesh: dolfinx.mesh.Mesh,
     meshtag: dolfinx.mesh.MeshTags,
-    comm: MPI.Intracomm,
+    comm: MPI.Comm,
     root: int,
 ) -> typing.Optional[dict[str, tuple[int, npt.NDArray]]]:
     """