feat: Add cutout schema version

.github/workflows/test.yaml

@@ -76,6 +76,12 @@ jobs:
       with:
         fetch-depth: 0 # Needed for setuptools_scm
 
+    - name: Determine cutout schema version
+      # Cached cutouts are only valid for the schema version they were created with
+      run: |
+        echo "CUTOUT_SCHEMA_VERSION=$(sed -n 's/^CUTOUT_SCHEMA_VERSION = \([0-9]*\).*/\1/p' atlite/schema.py)" >> $GITHUB_ENV
+      shell: bash
+
     - name: Set up Python ${{ matrix.python-version }} on ${{ matrix.os }}
       uses: actions/setup-python@v7
       with:
@@ -86,7 +92,7 @@ jobs:
       uses: actions/cache@v6
       with:
         path: ${{ env.CACHE_PATH }}
-        key: retrieved-cutouts-${{ env.today }}
+        key: retrieved-cutouts-schema${{ env.CUTOUT_SCHEMA_VERSION }}-${{ env.today }}
         enableCrossOsArchive: true
 
     - name: Download package

RELEASE_NOTES.rst

@@ -20,6 +20,13 @@ Upcoming Release
 **Features**
 
 * Add support for Python 3.14.
+* Feat: Cutouts data format (schema) is now versioned using a the attribute ``atlite_cutout_schema_version``.
+  Newly created cutouts carry the version recommended by the ``atlite`` version used to create them (version 1 at the moment)
+  Old cutouts are still supported but need to be migrated with
+  ``python -m atlite.migrate <path> --output <new path>`` or 
+  ``atlite.migrate.migrate_cutout(path, output=new_path)``.
+  Cutouts that have an incompatible data schema will in the future be rejected and raise a ``IncompatibleCutoutError``.
+  See :doc:`migrating_cutouts` for details (https://github.com/PyPSA/atlite/pull/533).
 
 **Bug fixes**
 

atlite/cutout.py

@@ -66,6 +66,13 @@
     compute_intersectionmatrix,
     get_coords,
 )
+from atlite.schema import (
+    CUTOUT_SCHEMA_VERSION,
+    SCHEMA_VERSION_ATTR,
+    check_schema_version,
+    read_schema_version,
+    schema_version_from_attrs,
+)
 from atlite.utils import CachedAttribute
 
 logger = logging.getLogger(__name__)
@@ -159,8 +166,16 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
             If required arguments are missing when building a new cutout.
         ValueError
             If ``bounds`` has an invalid format.
+        atlite.schema.IncompatibleCutoutError
+            If the cutout's schema version is not supported by this version of
+            atlite.
 
-        """
+        Warns
+        -----
+        atlite.schema.OutdatedCutoutWarning
+            If the cutout has an older, still supported schema version.
+
+        """  # noqa: DOC502 (IncompatibleCutoutError is raised by check_schema_version)
         path = Path(path).with_suffix(".nc")
         chunks = cutoutparams.pop("chunks", {"time": 100})
         if isinstance(chunks, dict):
@@ -171,6 +186,7 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
         # Three cases. First, cutout exists -> take the data.
         # Second, data is given -> take it. Third, else -> build a new cutout
         if path.is_file():
+            check_schema_version(read_schema_version(path), path)
             data = xr.open_dataset(str(path))
             data = data.chunk(chunks)
             data.attrs.update(storable_chunks)
@@ -182,6 +198,7 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
                 )
         elif "data" in cutoutparams:
             data = cutoutparams.pop("data")
+            check_schema_version(schema_version_from_attrs(data.attrs))
         else:
             logger.info("Building new cutout %s", path)
 
@@ -216,6 +233,7 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
             coords = get_coords(x, y, time, **cutoutparams)
 
             attrs = {
+                SCHEMA_VERSION_ATTR: CUTOUT_SCHEMA_VERSION,
                 "module": module,
                 "prepared_features": [],
                 **storable_chunks,
@@ -232,6 +250,11 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
         self.path = path
         self.data = data
 
+    @property
+    def schema_version(self) -> int:
+        """Schema version of the cutout's file format."""
+        return schema_version_from_attrs(self.data.attrs)
+
     @property
     def name(self) -> str:
         """Name of the cutout."""
@@ -434,9 +457,22 @@ def merge(
         merged : Cutout
             Merged cutout.
 
+        Raises
+        ------
+        ValueError
+            If the cutouts have different schema versions.
+
         """
         assert isinstance(other, Cutout)
 
+        if self.schema_version != other.schema_version:
+            raise ValueError(
+                f"Cannot merge cutouts with different schema versions "
+                f"({self.schema_version} and {other.schema_version}). "
+                "Migrate both to the current version first, then merge. "
+                "Use `python -m atlite.migrate` or `atlite.migrate.migrate_cutout`."
+            )
+
         if path is None:
             path = mktemp(
                 prefix=f"{self.path.stem}-",

atlite/data.py

@@ -22,6 +22,10 @@
 from numpy import atleast_1d
 
 from atlite.datasets import modules as datamodules
+from atlite.schema import (
+    SCHEMA_VERSION_ATTR,
+    require_current_schema_version,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Sequence
@@ -252,8 +256,13 @@ def cutout_prepare(
         If ``tmpdir`` is None.
     FileNotFoundError
         If ``tmpdir`` does not exist.
+    atlite.schema.IncompatibleCutoutError
+        If features is to be added to an existing cutout that is not the recommended
+        schema version of the installed versino of ``atlite``.
+    RuntimeError
+        If the data of a module sets the cutout's schema version attribute.
 
-    """
+    """  # noqa: DOC502 (IncompatibleCutoutError is raised by require_current_schema_version)
     if dask_kwargs is None:
         dask_kwargs = {}
 
@@ -264,6 +273,8 @@ def cutout_prepare(
         logger.info("Cutout already prepared.")
         return cutout
 
+    require_current_schema_version(cutout.schema_version, cutout.path)
+
     if tmpdir is None:
         raise ValueError("tmpdir cannot be None")
     temp_dir_path = Path(tmpdir)
@@ -296,6 +307,10 @@ def cutout_prepare(
             monthly_requests=monthly_requests,
             concurrent_requests=concurrent_requests,
         )
+        if SCHEMA_VERSION_ATTR in ds.attrs:
+            raise RuntimeError(
+                f"Data of module '{module}' must not contain attribute '{SCHEMA_VERSION_ATTR}'."
+            )
         prepared |= set(missing_features)
 
         cutout.data.attrs.update({"prepared_features": list(prepared)})

atlite/migrate.py

@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: Contributors to atlite <https://github.com/pypsa/atlite>
+#
+# SPDX-License-Identifier: MIT
+"""
+Explicit migration of cutouts to the current schema version.
+
+Write the migrated cutout to a new file and keep the original::
+
+    python -m atlite.migrate path/to/cutout.nc --output path/to/migrated.nc
+
+or, from Python::
+
+    from atlite.migrate import migrate_cutout
+
+    migrate_cutout("path/to/cutout.nc", output="path/to/migrated.nc")
+
+Overwriting the original file requires ``--in-place`` (``in_place=True``).
+
+Each migration step upgrades a cutout file by exactly one schema version and is
+registered in :data:`MIGRATIONS` under the version it upgrades from.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import shutil
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import xarray as xr
+
+from atlite.schema import (
+    CUTOUT_SCHEMA_VERSION,
+    SCHEMA_VERSION_ATTR,
+    IncompatibleCutoutError,
+    read_schema_version,
+    require_supported_schema_version,
+    schema_version_from_attrs,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from atlite._types import PathLike
+
+logger = logging.getLogger(__name__)
+
+
+def _migrate_v0_to_v1(path: Path) -> None:
+    """Migrate a cutout file in place from schema version 0 to 1."""
+    # v1 only adds the version attribute
+    # Appending a dataset without variables updates the meta attributes in the header and leaves the data untouched
+    xr.Dataset(attrs={SCHEMA_VERSION_ATTR: 1}).to_netcdf(path, mode="a")
+
+
+# Maps the schema version a step upgrades from to the step itself
+MIGRATIONS: dict[int, Callable[[Path], None]] = {
+    0: _migrate_v0_to_v1,
+}
+
+
+def migrate_cutout(
+    path: PathLike, output: PathLike | None = None, *, in_place: bool = False
+) -> None:
+    """
+    Migrate a cutout file to the current schema version.
+
+    Exactly one of ``output`` and ``in_place=True`` must be given.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path of the cutout to migrate.
+    output : str or Path, optional
+        Path where to write the migrated cutout to. Must not exist yet.
+    in_place : bool, default False
+        Overwrite the original cutout with the migrated cutout.
+
+    Raises
+    ------
+    ValueError
+        If neither or both of ``output`` and ``in_place=True`` are given.
+    IncompatibleCutoutError
+        If the file is not an atlite cutout or its schema version cannot be
+        migrated by this atlite version.
+    FileExistsError
+        If ``output`` already exists.
+
+    """
+    # Overwriting the original must be an explicit choice
+    if (output is None) == (not in_place):
+        raise ValueError("Specify either an output file or in_place=True")
+
+    path = Path(path)
+    with xr.open_dataset(path) as ds:
+        attrs = dict(ds.attrs)
+    # Files without a version attribute are treated as v0 cutouts, so make sure
+    # it is a cutout at all before stamping a version on it
+    if "module" not in attrs:
+        raise IncompatibleCutoutError(
+            f"'{path}' does not appear to be an atlite cutout and is missing the attribute 'module'. "
+            f"Aborting."
+        )
+    version = schema_version_from_attrs(attrs)
+    require_supported_schema_version(version, path)
+
+    if output is None:
+        target = path
+    else:
+        output = Path(output)
+        if output.exists():
+            raise FileExistsError(f"Output file '{output}' already exists.")
+        # Migrate a hidden copy next to the output, so `output` only appears
+        # once it is complete and the original is not touched before, avoiding
+        # corruption on failed upgrades
+        target = output.with_name(f".{output.name}.migration")
+        shutil.copyfile(path, target)
+
+    try:
+        # Actual migration loop. Call all migration functions in order
+        for current in range(version, CUTOUT_SCHEMA_VERSION):
+            logger.info(
+                "Migrating '%s' from schema version %d to %d.",
+                path,
+                current,
+                current + 1,
+            )
+            MIGRATIONS[current](target)
+    except BaseException:
+        # Also on KeyboardInterrupt: don't leave a half-migrated copy behind
+        if target != path:
+            target.unlink(missing_ok=True)
+        raise
+
+    if output is not None:
+        target.rename(output)
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """
+    Migrate a cutout from the command line.
+
+    Parameters
+    ----------
+    argv : sequence of str, optional
+        Command line arguments, defaults to ``sys.argv[1:]``.
+
+    Returns
+    -------
+    int
+        Exit code.
+
+    """
+    parser = argparse.ArgumentParser(
+        prog="python -m atlite.migrate",
+        description="Migrate an atlite cutout to the current schema version "
+        f"({CUTOUT_SCHEMA_VERSION}).",
+    )
+    parser.add_argument("cutout", type=Path, help="Path to the cutout file.")
+    # Mirrors migrate_cutout: exactly one destination is required
+    destination = parser.add_mutually_exclusive_group(required=True)
+    destination.add_argument(
+        "-o", "--output", type=Path, help="Write to this file, keep the original."
+    )
+    destination.add_argument(
+        "--in-place", action="store_true", help="Overwrite the original file."
+    )
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+
+    try:
+        version = read_schema_version(args.cutout)
+        migrate_cutout(args.cutout, args.output, in_place=args.in_place)
+    except (IncompatibleCutoutError, OSError) as err:
+        # Expected failures: report them without a traceback
+        logger.error("%s", err)
+        return 1
+
+    logger.info(
+        "'%s' has schema version %d (was %d).",
+        args.output or args.cutout,
+        CUTOUT_SCHEMA_VERSION,
+        version,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())