WIP: Harden types

.pre-commit-config.yaml

@@ -24,7 +24,6 @@ repos:
   - id: ruff-check
     exclude: '(dev/.*|.*_)\.py$'
     args:
-    - --line-length=120
     - --fix
     - --exit-non-zero-on-fix
   - id: ruff-format

pixi.toml

@@ -23,6 +23,9 @@ pytest-cov = "*"
 pytest-mock = "*"
 ruff = "*"
 
+[feature.dev.pypi-dependencies]
+pandas-stubs = "*"
+
 [feature.docs.dependencies]
 nbsphinx = ">=0.9.8,<0.10"
 pandoc = "*"
@@ -40,7 +43,7 @@ earthkit-plots = ">=0.5.0"
 [tasks]
 qa = "pre-commit run --all-files"
 template-update = "pre-commit run --all-files cruft -c .pre-commit-config-cruft.yaml"
-type-check = "python -m mypy . --no-namespace-packages"
+type-check = "MYPYPATH=src python -m mypy src tests"
 unit-tests = "python -m pytest -vv --cov=. --cov-report=html --doctest-glob='*.md' --doctest-glob='*.rst'"
 
 [workspace]

pyproject.toml

@@ -35,6 +35,7 @@ optional-dependencies.dev = [
   "ipykernel",
   "ipython",
   "mypy",
+  "pandas-stubs",
   "pre-commit",
   "pytest",
   "pytest-cov",
@@ -52,6 +53,24 @@ optional-dependencies.docs = [
 [tool.coverage.run]
 branch = true
 
+[tool.mypy]
+exclude = ["build", "dist"]
+explicit_package_bases = true
+files = ["src", "tests"]
+namespace_packages = true
+
+[[tool.mypy.overrides]]
+ignore_missing_imports = true
+module = [
+  "earthkit",
+  "earthkit.data",
+  "earthkit.data.*",
+  "xclim",
+  "xclim.*",
+  "xsdba",
+  "xsdba.*"
+]
+
 [tool.pytest.ini_options]
 addopts = "-vv --cov=. --cov-report=html --doctest-glob='*.md' --doctest-glob='*.rst'"
 
@@ -86,14 +105,17 @@ lint.per-file-ignores."docs/conf.py" = [
 [tool.setuptools.dynamic]
 readme = {file = ["README.md"], content-type = "text/markdown"}
 
+[tool.setuptools.package-data]
+"earthkit.climate" = ["py.typed"]
+
 # ---- Namespace package configuration ----
 [tool.setuptools.packages.find]
 where = ["src"]
 
 [tool.setuptools_scm]
 local_scheme = "no-local-version"
-write_to = "src/earthkit/climate/version.py"
-write_to_template = '''
+version_file = "src/earthkit/climate/version.py"
+version_file_template = '''
 # Do not change! Do not track in version control!
 __version__ = "{version}"
 '''

src/earthkit/climate/__init__.py

@@ -9,13 +9,13 @@
 try:
     # NOTE: the `version.py` file must not be present in the git repository
     #   as it is generated by setuptools at install time
-    from .version import __version__
+    from earthkit.climate.version import __version__
 except ImportError:  # pragma: no cover
     # Local copy or not installed with setuptools
     __version__ = "999"
 
 # Avoid importing optional heavy submodules at package import time to keep
 # test collection lightweight and not require optional dependencies (e.g., xclim).
-from .utils import conversions
+from earthkit.climate.utils import conversions
 
-__all__ = [conversions, __version__]
+__all__ = ["conversions", "__version__"]

src/earthkit/climate/api/wrapper.py

@@ -6,15 +6,30 @@
 # granted to it by virtue of its status as an intergovernmental organisation nor
 # does it submit to any jurisdiction.
 
+from __future__ import annotations
+
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable, Dict, Union
+from typing import Any, Protocol, TypeAlias
 
 import xarray as xr
 
 from earthkit.climate.utils import conversions, provenance, units
 
+IndicatorInput: TypeAlias = conversions.EarthkitData | xr.Dataset
+
+
+class XclimIndicator(Protocol):
+    """Protocol for xclim indicators wrapped by this module."""
+
+    parameters: Any
+    cf_attrs: Any
+    compute: Any
+
+    def __call__(self, *args, **kwargs) -> xr.Dataset | xr.DataArray: ...
+
 
-def wrap_xclim_indicator(xclim_fn: Callable) -> Callable:
+def wrap_xclim_indicator(xclim_fn: XclimIndicator) -> Callable[..., conversions.EarthkitData]:
     """
     Wraps an xclim indicator to handle Earthkit inputs and unit alignment.
 
@@ -31,7 +46,7 @@ def wrap_xclim_indicator(xclim_fn: Callable) -> Callable:
 
     @wraps(xclim_fn)
     def wrapper(
-        earthkit_input: Union[conversions.EarthkitData, xr.Dataset],
+        earthkit_input: IndicatorInput,
         *args,
         **kwargs,
     ) -> conversions.EarthkitData:
@@ -52,7 +67,7 @@ def wrapper(
         conversions.EarthkitData
             The result of the indicator calculation wrapped as an Earthkit object.
         """
-        metadata: Dict[str, Any] = {}
+        metadata: conversions.MetadataDict = {}
 
         # --- STEP 1: Load & Standardize Main Data ---
         # Convert Earthkit object to xarray Dataset
@@ -67,7 +82,7 @@ def wrapper(
 
         # --- STEP 2: Execution ---
         # We pass the single merged dataset (ds) and the variable name mappings
-        output_dataset: xr.Dataset = xclim_fn(ds=dataset, *args, **kwargs)
+        output_dataset = xclim_fn(ds=dataset, *args, **kwargs)
 
         # --- STEP 3: Provenance & Output ---
         metadata = provenance.add_indicator_provenance(metadata, xclim_fn, dataset, **kwargs)

src/earthkit/climate/indicators/__init__.py

@@ -12,3 +12,8 @@
 the input and include attributes for CF metadata (cell methods), references,
 keywords, and more.
 """
+
+import earthkit.climate.indicators.precipitation as precipitation
+import earthkit.climate.indicators.temperature as temperature
+
+__all__ = ["precipitation", "temperature"]

src/earthkit/climate/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561 typed package support.

src/earthkit/climate/utils/conversions.py

@@ -10,19 +10,20 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, Mapping, Tuple
+from collections.abc import Mapping
+from typing import Any, TypeAlias
 
 import earthkit.data as ekd
 import xarray
 
-EarthkitData = ekd.FieldList | ekd.Field
-MetadataDict = Dict[str, Any]
+EarthkitData: TypeAlias = ekd.FieldList | ekd.Field
+MetadataDict: TypeAlias = dict[str, Any]
 
 
 def to_xarray_dataset(
     earthkit_input: EarthkitData | xarray.Dataset,
     metadata: Mapping[str, Any] | None = None,
-) -> Tuple[xarray.Dataset, MetadataDict]:
+) -> tuple[xarray.Dataset, MetadataDict]:
     """
     Convert Earthkit-like data to an ``xarray.Dataset`` and gather metadata.
 

src/earthkit/climate/utils/provenance.py

@@ -11,16 +11,24 @@
 from __future__ import annotations
 
 import inspect
-from typing import Any
+from typing import Any, Protocol
 
 import xarray as xr
 
 from earthkit.climate.utils.conversions import MetadataDict
 
 
+class IndicatorWithCompute(Protocol):
+    """Indicator-like object with a callable interface and a compute method."""
+
+    def compute(self, *args: Any, **kwargs: Any) -> object: ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> object: ...
+
+
 def add_indicator_provenance(
     metadata: MetadataDict,
-    indicator: Any,
+    indicator: IndicatorWithCompute,
     dataset: xr.Dataset,
     **kwargs: Any,
 ) -> MetadataDict:

tests/test_00_version.py

@@ -1,5 +1,11 @@
-from earthkit import climate
+from pathlib import Path
+
+import earthkit.climate as climate
 
 
 def test_version() -> None:
     assert climate.__version__ != "999"
+
+
+def test_py_typed_marker_present() -> None:
+    assert Path(climate.__file__).with_name("py.typed").is_file()