General cleanup (#35)

* 🔧 config: reorganize dependency groups and update package versions
* 💄 ui: change typing_extensions.Doc for docstrings
* 🔧 config(ci): nox and py13
* 👷 build: CI pass step

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

Author: nstarman

.github/workflows/ci.yml

@@ -35,14 +35,14 @@ jobs:
           echo "::add-matcher::$GITHUB_WORKSPACE/.github/matchers/pylint.json"
           uv run nox -s pylint
 
-  checks:
+  tests:
     name: Check Python ${{ matrix.python-version }} on ${{ matrix.runs-on }}
     runs-on: ${{ matrix.runs-on }}
     needs: [pre-commit]
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.12"]
+        python-version: ["3.10", "3.13"]
         runs-on: [ubuntu-latest, macos-latest, windows-latest]
 
     steps:
@@ -61,3 +61,14 @@ jobs:
         uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
+
+  status:
+    name: CI Pass
+    runs-on: ubuntu-latest
+    needs: [pre-commit, tests]
+    if: always()
+    steps:
+      - name: All required jobs passed
+        uses: re-actors/alls-green@release/v1
+        with:
+          jobs: ${{ toJSON(needs) }}

noxfile.py

@@ -9,83 +9,96 @@
 
 nox.needs_version = ">=2024.3.2"
 nox.options.sessions = [
-    "tests",
     # Linting
     "lint",
     "pylint",
+    "precommit",
     # Testing
-    "test",
+    "tests",
     # Packaging
     "build",
 ]
-nox.options.default_venv_backend = "uv|virtualenv"
-
-
-@nox.session
-def check(session: nox.Session, /) -> None:
-    """Run all checks."""
-    lint(session)
-    test(session)
+nox.options.default_venv_backend = "uv"
 
 
 # =============================================================================
+# Linting
 
 
-@nox.session
+@nox.session(venv_backend="uv")
 def lint(session: nox.Session, /) -> None:
     """Run the linter."""
-    precommit(session)
-    pylint(session)
+    precommit(session)  # reuse pre-commit session
+    pylint(session)  # reuse pylint session
 
 
-@nox.session
+@nox.session(venv_backend="uv")
 def precommit(session: nox.Session, /) -> None:
-    """Run the pre-commit hooks."""
-    session.run(
+    """Run pre-commit."""
+    session.run_install(
         "uv",
-        "run",
-        "pre-commit",
-        "run",
-        "--all-files",
-        "--show-diff-on-failure",
-        *session.posargs,
+        "sync",
+        "--group=lint",
+        f"--python={session.virtualenv.location}",
+        env={"UV_PROJECT_ENVIRONMENT": session.virtualenv.location},
     )
+    session.run("pre-commit", "run", "--all-files", *session.posargs)
 
 
-@nox.session
+@nox.session(venv_backend="uv")
 def pylint(session: nox.Session, /) -> None:
     """Run PyLint."""
-    # This needs to be installed into the package environment, and is slower
-    # than a pre-commit check
-    session.run("uv", "sync", "--group", "pylint")
-    session.run("uv", "run", "pylint", "xmmutablemap", *session.posargs)
+    session.run_install(
+        "uv",
+        "sync",
+        "--group=lint",
+        f"--python={session.virtualenv.location}",
+        env={"UV_PROJECT_ENVIRONMENT": session.virtualenv.location},
+    )
+    session.run("pylint", "xmmutablemap", *session.posargs)
 
 
 # =============================================================================
 # Testing
 
 
-@nox.session
-def test(session: nox.Session, /) -> None:
-    """Run the tests."""
-    session.run("uv", "sync", "--group", "test")
-    session.run("uv", "run", "pytest", *session.posargs)
+@nox.session(venv_backend="uv")
+def tests(session: nox.Session, /) -> None:
+    """Run the unit and regular tests."""
+    session.run_install(
+        "uv",
+        "sync",
+        "--group=test",
+        f"--python={session.virtualenv.location}",
+        env={"UV_PROJECT_ENVIRONMENT": session.virtualenv.location},
+    )
+    session.run("pytest", *session.posargs)
 
 
 # =============================================================================
 # Packaging
 
 
-@nox.session
+@nox.session(venv_backend="uv")
 def rm_build(_: nox.Session, /) -> None:
     """Remove the build directory."""
     build_path = DIR.joinpath("build")
     if build_path.exists():
         shutil.rmtree(build_path)
 
 
-@nox.session
+@nox.session(venv_backend="uv")
 def build(session: nox.Session, /) -> None:
     """Build an SDist and wheel."""
-    rm_build(session)
-    session.run("uv", "build", *session.posargs)
+    build_path = DIR.joinpath("build")
+    if build_path.exists():
+        shutil.rmtree(build_path)
+
+    session.run_install(
+        "uv",
+        "sync",
+        "--group=build",
+        f"--python={session.virtualenv.location}",
+        env={"UV_PROJECT_ENVIRONMENT": session.virtualenv.location},
+    )
+    session.run("build")

pyproject.toml

@@ -36,31 +36,34 @@ Homepage = "https://github.com/GalacticDynamics/xmmutablemap"
 Discussions = "https://github.com/GalacticDynamics/xmmutablemap/discussions"
 Changelog = "https://github.com/GalacticDynamics/xmmutablemap/releases"
 
+
+[build-system]
+requires = ["hatchling", "hatch-vcs"]
+build-backend = "hatchling.build"
+
+
 [dependency-groups]
 dev = [
   "ipykernel>=6.29.5",
   "cz-conventional-gitmoji>=0.6.1",
-  "pre-commit>=4.1.0",
+  {include-group = "lint"},
   {include-group = "nox"},
   {include-group = "test"},
 ]
-pylint = [
+lint = [
+    "pre-commit>=4.1.0",
     "pylint>=3.3.8",
 ]
 nox = [
     "nox>=2024.10.9",
 ]
 test = [
-  "pytest>=8.3.3",
-  "pytest-cov>=3",
+  "pytest>=8.4.2",
+  "pytest-cov>=6.2.1",
   "pytest-github-actions-annotate-failures>=0.2.0",
-  "sybil>=8.0.0",
+  "sybil[pytest]>=9.2.0",
 ]
 
-[build-system]
-requires = ["hatchling", "hatch-vcs"]
-build-backend = "hatchling.build"
-
 
 [tool.hatch]
 build.hooks.vcs.version-file = "src/xmmutablemap/_version.py"
@@ -72,7 +75,7 @@ name = "cz_gitmoji"
 
 
 [tool.pytest.ini_options]
-minversion = "8.3"
+minversion = "8"
 addopts = [
   "--showlocals",
   "--strict-config",

src/xmmutablemap/_core.py

@@ -6,8 +6,7 @@
 __all__ = ("ImmutableMap",)
 
 from collections.abc import ItemsView, Iterable, Iterator, KeysView, Mapping, ValuesView
-from typing import Annotated, Any, TypeVar, overload
-from typing_extensions import Doc
+from typing import Any, TypeVar, overload
 
 from jax.tree_util import register_pytree_node_class
 
@@ -234,16 +233,17 @@ def __repr__(self) -> str:
     # ===========================================
     # JAX PyTree
 
-    def tree_flatten(
-        self,
-    ) -> tuple[
-        Annotated[tuple[V, ...], Doc("The values.")],
-        Annotated[tuple[K, ...], Doc("The keys as auxiliary data.")],
-    ]:
+    def tree_flatten(self) -> tuple[tuple[V, ...], tuple[K, ...]]:
         """Flatten dict to the values (and keys).
 
         This is used for JAX's tree flattening.
 
+        Returns
+        -------
+        tuple[tuple[V, ...], tuple[K, ...]]
+            A tuple of (values, keys).
+            The keys are treated as auxiliary data.
+
         Examples
         --------
         >>> import jax
@@ -260,14 +260,19 @@ def tree_flatten(
 
     @classmethod
     def tree_unflatten(
-        cls,
-        aux_data: Annotated[tuple[K, ...], Doc("The keys.")],
-        children: Annotated[tuple[V, ...], Doc("The values.")],
+        cls, aux_data: tuple[K, ...], children: tuple[V, ...]
     ) -> "ImmutableMap[K, V]":
         """Unflatten into an ImmutableMap from the keys and values.
 
         This is used for JAX's tree un-flattening.
 
+        Parameters
+        ----------
+        aux_data : tuple[K, ...]
+            The keys.
+        children : tuple[V, ...]
+            The values.
+
         Examples
         --------
         >>> import jax