feat(repositories): Enable working with eu stream

.github/workflows/branch_ci.yml

@@ -17,7 +17,7 @@ jobs:
     with:
       enable_linting: true
       enable_typechecking: true
-      tests_folder: "src/nwp_consumer"
+      tests_folder: "src/nwp_consumer -t src"
       tests_runner: "unittest"
       containerfile: Containerfile
 

.gitignore

@@ -52,7 +52,6 @@ docs/
 
 # Environments
 .venv
-uv.lock
 .env
 
 # mypy

Makefile

@@ -0,0 +1,14 @@
+.PHONY: test lint typecheck
+
+test:
+	uv run python -m unittest discover -s src/nwp_consumer -t src -p "test_*.py"
+
+test-cov:
+	uv run python -m xmlrunner discover -s src/nwp_consumer -t src -p "test_*.py"
+
+lint:
+	uv run ruff check --fix .
+
+typecheck:
+	uv run ty check src
+

README.md

@@ -108,15 +108,15 @@ repository.
 
 ### Linting and static type checking
 
-This project uses [MyPy](https://mypy.readthedocs.io/en/stable/) for static type checking
+This project uses [ty](https://github.com/astral-sh/ty) for static type checking
 and [Ruff](https://docs.astral.sh/ruff/) for linting.
 Installing the development dependencies makes them available in your virtual environment.
 
-Use them via:
+Use them via the provided Makefile targets:
 
 ```bash
-$ python -m mypy .
-$ python -m ruff check .
+$ make typecheck
+$ make lint
 ```
 
 Be sure to do this periodically while developing to catch any errors early
@@ -128,7 +128,7 @@ but it prevents accidental creation of a whole suite of bugs.
 Run the unittests with:
 
 ```bash
-$ python -m unittest discover -s src/nwp_consumer -p "test_*.py"
+$ make test
 ```
 
 ## Further reading

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
     "s3fs == 2024.9.0",
     "xarray == 2025.1.0",
     "zarr == 2.18.3",
-    "pandas < 3.0"
+    "pandas < 3.0",
 ]
 
 [dependency-groups]
@@ -44,8 +44,7 @@ dev = [
     "unittest-xml-reporting == 3.2.0",
     "hypothesis == 6.115.3",
     # Linting
-    "returns[compatible-mypy]",
-    "ruff == 0.6.9",
+    "ruff==0.6.9",
     "pandas-stubs",
     "types-psutil",
     "types-pytz",
@@ -55,8 +54,8 @@ dev = [
     "pdoc >= 15.0.0",
     # IDE support
     "python-lsp-server",
-    "pylsp-mypy",
     "python-lsp-ruff",
+    "ty>=0.0.75",
 ]
 
 [project.scripts]
@@ -82,34 +81,6 @@ include = ["src"]
 exclude = ["**/__pycache__"]
 pythonVersion = "3.12"
 
-# MyPy configuration
-# * See https://mypy.readthedocs.io/en/stable/index.html
-[tool.mypy]
-python_version = "3.12"
-strict = true
-warn_unreachable = true
-warn_return_any = true
-disallow_untyped_defs = true
-plugins = [
-    "returns.contrib.mypy.returns_plugin",
-    "numpy.typing.mypy_plugin",
-]
-
-[[tool.mypy.overrides]]
-# Ignore missing imports for libraries that don't have them.
-# If they are ever made, remove from here!
-module = [
-    "ecmwfapi",
-    "cfgrib",
-    "botocore.session",
-    "botocore.client",
-    "joblib",
-    "ocf_blosc2",
-    "s3fs",
-    "zarr",
-]
-ignore_missing_imports = true
-
 # Ruff configuration
 # * See https://beta.ruff.rs/docs/
 [tool.ruff]
@@ -119,7 +90,6 @@ exclude = ["__init__.py"]
 
 [tool.ruff.lint]
 fixable = ["ALL"]
-ignore = ["ANN101", "ANN102"]
 select = [
     "F",   # pyflakes
     "E",   # pycodestyle

src/nwp_consumer/internal/entities/coordinates.py

@@ -370,7 +370,7 @@ def to_pandas(self) -> dict[str, pd.Index]:
             `NWPDimensionCoordinateMap.from_pandas` for the reverse operation.
 
         """
-        out_dict: dict[str, pd.Index] = {  # type: ignore
+        out_dict: dict[str, pd.Index] = {
             "init_time": pd.Index(
                 [
                     np.datetime64(t.astimezone(dt.UTC).replace(tzinfo=None), "ns")
@@ -556,7 +556,7 @@ def as_zeroed_dataarray(self, name: str, chunks: dict[str, int]) -> xr.DataArray
         """
         # Create a dask array of zeros with the shape of the dataset
         # * The values of this are ignored, only the shape and chunks are used
-        dummy_values = dask.array.zeros(  # type: ignore
+        dummy_values = dask.array.zeros(
             shape=list(self.shapemap.values()),
             chunks=tuple([chunks[k] for k in self.shapemap]),
         )

src/nwp_consumer/internal/entities/performance.py

@@ -66,7 +66,7 @@ def get_runtime(self) -> int:
 
     def run(self) -> None:
         """Run the thread."""
-        memory_start, cpu_start = self.get_usage()
+        memory_start, _ = self.get_usage()
         while not self.stop:
             new_memory, new_cpu = self.get_usage()
             # Memory is just a total, so get the delta

src/nwp_consumer/internal/entities/tensorstore.py

@@ -20,9 +20,12 @@
 from typing import Any
 
 import numpy as np
+import numpy.typing
 import pandas as pd
 import xarray as xr
 import zarr
+import zarr.errors
+import zarr.storage
 from returns.result import Failure, ResultE, Success
 
 from .coordinates import NWPDimensionCoordinateMap
@@ -141,7 +144,7 @@ def initialize_empty_store(
         try:
             if zarrdir.startswith("s3"):
                 store_result = cls._create_zarrstore_s3(zarrdir, filename)
-                store, path = store_result.unwrap()  # Can do this as exceptions are caught
+                store, path = store_result.unwrap()  # ty: ignore[invalid-assignment]
             else:
                 path = pathlib.Path("/".join((zarrdir, filename))).expanduser().as_posix()
                 store = zarr.storage.DirectoryStore(path)
@@ -174,9 +177,9 @@ def initialize_empty_store(
             )
             log.info("Created blank zarr store at '%s'", path)
             # Ensure the store is readable
-            store_da = xr.open_dataarray(store, engine="zarr")
+            store_da = xr.open_dataarray(store, engine="zarr")  # ty: ignore[invalid-argument-type]
         except zarr.errors.ContainsGroupError:
-            store_da = xr.open_dataarray(store, engine="zarr")
+            store_da = xr.open_dataarray(store, engine="zarr")  # ty: ignore[invalid-argument-type]
             if store_da.name != da.name:  # TODO: Also check for equality of coordinates
                 return Failure(
                     OSError(
@@ -540,7 +543,7 @@ def update_attrs(self, attrs: dict[str, str]) -> ResultE[str]:
         """
         group: zarr.Group = zarr.open_group(self.path)
         group.attrs.update(attrs)
-        zarr.consolidate_metadata(self.path)
+        zarr.consolidate_metadata(self.path)  # ty: ignore[invalid-argument-type]
         return Success(self.path)
 
     def missing_times(self) -> ResultE[list[dt.datetime]]:
@@ -577,7 +580,7 @@ def missing_times(self) -> ResultE[list[dt.datetime]]:
         return Success(missing_times)
 
     @staticmethod
-    def _create_zarrstore_s3(s3_folder: str, filename: str) -> ResultE[tuple[MutableMapping, str]]:  # type: ignore
+    def _create_zarrstore_s3(s3_folder: str, filename: str) -> ResultE[tuple[MutableMapping, str]]:
         """Create a mutable mapping to an S3 store.
 
         Authentication with S3 is done via botocore's credential discovery.
@@ -590,6 +593,7 @@ def _create_zarrstore_s3(s3_folder: str, filename: str) -> ResultE[tuple[Mutable
           - https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#configuring-credentials
         """
         import s3fs
+        import s3fs.mapping
 
         if not s3_folder.startswith("s3://"):
             return Failure(

src/nwp_consumer/internal/repositories/raw_repositories/ceda.py

@@ -87,9 +87,9 @@
 from collections.abc import Callable, Iterator
 from typing import override
 
-import cfgrib
 import numpy as np
 import xarray as xr
+from cfgrib.xarray_store import open_datasets
 from joblib import delayed
 from returns.result import Failure, ResultE, Success
 
@@ -143,7 +143,7 @@ def repository() -> entities.RawRepositoryMetadata:
     def model() -> entities.ModelMetadata:
         requested_model: str = os.getenv("MODEL", default="default")
         if requested_model not in CEDARawRepository.repository().available_models:
-            log.warn(
+            log.warning(
                 f"Unknown model '{requested_model}' requested, falling back to default ",
                 "CEDA repository only supports "
                 f"'{list(CEDARawRepository.repository().available_models.keys())}'. "
@@ -309,7 +309,7 @@ def _convert_global(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
             path: The path to the file to convert.
         """
         try:
-            dss: list[xr.Dataset] = cfgrib.open_datasets(path)
+            dss: list[xr.Dataset] = open_datasets(path)
         except Exception as e:
             return Failure(
                 OSError(
@@ -383,7 +383,7 @@ def _convert_ukv(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
         # * cfgrib loads multiple hypercubes for a single multi-parameter grib file
         # * Can also set backend_kwargs={"indexpath": ""}, to avoid the index file
         try:
-            dss: list[xr.Dataset] = cfgrib.open_datasets(
+            dss: list[xr.Dataset] = open_datasets(
                 path=path.as_posix(),
                 chunks={"time": 1, "step": -1, "variable": -1, "x": "auto", "y": "auto"},
                 backend_kwargs={"indexpath": ""},

src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_mars.py

@@ -17,8 +17,8 @@
 from collections.abc import Callable, Iterator
 from typing import override
 
-import cfgrib
 import xarray as xr
+from cfgrib.xarray_store import open_datasets
 from ecmwfapi import ECMWFService
 from joblib import delayed
 from returns.result import Failure, ResultE, Success
@@ -239,7 +239,7 @@ def repository() -> entities.RawRepositoryMetadata:
     def model() -> entities.ModelMetadata:
         requested_model: str = os.getenv("MODEL", default="default")
         if requested_model not in ECMWFMARSRawRepository.repository().available_models:
-            log.warn(
+            log.warning(
                 f"Unknown model '{requested_model}' requested, falling back to default ",
                 "ECMWF-MARS repository only supports "
                 f"'{list(ECMWFMARSRawRepository.repository().available_models.keys())}'. "
@@ -325,7 +325,7 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
             path: The path to the file to convert.
         """
         try:
-            dss: list[xr.Dataset] = cfgrib.open_datasets(
+            dss: list[xr.Dataset] = open_datasets(
                 path=path.as_posix(),
                 chunks={"time": 1, "step": -1, "longitude": -1, "latitude": -1},
                 backend_kwargs={"indexpath": ""},

src/nwp_consumer/internal/repositories/raw_repositories/ecmwf_realtime.py

@@ -21,9 +21,10 @@
 time of 2024-10-25 00:00 and a target time of 2024-10-26 01:00 (step of 25 hours).
 
 The file contents is specific to the order agreed with the data provider.
-For the order that OCF has created, there are four distinct datasets.
-This is because OCF has ordered two separate regions and 17 variables,
-which are split across two datasets.
+OCF orders a single west-europe region alongside a separate india region,
+and the variables are split across multiple datasets within each file.
+Datasets that overlap the requested model region are cropped to it, which
+allows building a sub-region store (e.g. nl) from the larger west-europe order.
 
 Also, some of the data contains larger steps than we are interested in due
 to necessities in the order creation process.
@@ -38,10 +39,10 @@
 from collections.abc import Callable, Iterator
 from typing import override
 
-import cfgrib
 import numpy as np
 import s3fs
 import xarray as xr
+from cfgrib.xarray_store import open_datasets
 from joblib import delayed
 from returns.result import Failure, ResultE, Success
 
@@ -91,6 +92,9 @@ def repository() -> entities.RawRepositoryMetadata:
                 "hres-ifs-nl": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region(
                     "nl",
                 ).with_max_step(84).with_delay_minutes(60 * 7),
+                "hres-ifs-west-europe": entities.Models.ECMWF_HRES_IFS_0P1DEGREE.with_region(
+                    "west-europe",
+                ).with_delay_minutes(60 * 7),
             },
         )
 
@@ -252,7 +256,7 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
             path: The path to the grib file.
         """
         try:
-            dss: list[xr.Dataset] = cfgrib.open_datasets(path.as_posix())
+            dss: list[xr.Dataset] = open_datasets(path.as_posix())
         except Exception as e:
             return Failure(
                 OSError(
@@ -273,14 +277,22 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
         expected_lats = ECMWFRealTimeS3RawRepository.model().expected_coordinates.latitude
         expected_steps = ECMWFRealTimeS3RawRepository.model().expected_coordinates.step
 
+        north, west, south, east = \
+            ECMWFRealTimeS3RawRepository.model().expected_coordinates.nwse()
+
         for i, ds in enumerate(dss):
-            # ECMWF Realtime provides all regions in one set of datasets,
-            # so distinguish via their coordinates
-            step = np.timedelta64(ds.coords["step"].values, "h").astype(int) # type: ignore[arg-type]
+            # ECMWF Realtime provides multiple orders in one set of datasets
+            # (e.g. west-europe and india), so distinguish via their coordinates.
+            # Datasets that overlap the requested region are kept and cropped to it,
+            # which allows a sub-region store (e.g. nl) to be built from the
+            # larger west-europe order.
+            ds_lons = ds.coords["longitude"].values
+            ds_lats = ds.coords["latitude"].values
+            step = np.timedelta64(ds.coords["step"].values, "h").astype(int)  # ty: ignore[invalid-argument-type]
             is_relevant_dataset_predicate: bool = (
                 (expected_lons is not None and expected_lats is not None)
-                and (expected_lons[0] <= max(ds.coords["longitude"].values) <= expected_lons[-1])
-                and (expected_lats[-1] <= max(ds.coords["latitude"].values) <= expected_lats[0])
+                and (min(ds_lons) <= expected_lons[-1] and max(ds_lons) >= expected_lons[0])
+                and (min(ds_lats) <= expected_lats[0] and max(ds_lats) >= expected_lats[-1])
                 and (expected_steps[0] <= step <= expected_steps[-1])
             )
             if not is_relevant_dataset_predicate:
@@ -311,6 +323,15 @@ def _convert(path: pathlib.Path) -> ResultE[list[xr.DataArray]]:
                     .sortby(variables=["step", "variable", "longitude"])
                     .sortby(variables="latitude", ascending=False)
                 )
+                # Crop to the requested model region, since the order may cover
+                # a larger extent (e.g. cropping nl from the west-europe order).
+                # Round coordinates to the store's 4 d.p. precision first, else
+                # floating-point grid-edge differences drop boundary points and
+                # misalign the write with chunk boundaries.
+                da = da.assign_coords(
+                    latitude=[float(f"{v:.4f}") for v in da["latitude"].values],
+                    longitude=[float(f"{v:.4f}") for v in da["longitude"].values],
+                ).sel(latitude=slice(north, south), longitude=slice(west, east))
 
             except Exception as e:
                 return Failure(