Merge pull request #2 from brews/drop_dtr_add_types

Drop DtrRun, add types, add mypy type checking in CI

.github/workflows/test.yaml

@@ -2,12 +2,14 @@ name: Test
 
 on:
   push:
-    branches: "main"
+    branches:
+      - "main"
   pull_request:
-    branches: "main"
+    branches:
+      - "main"
 
 jobs:
-  build:
+  test:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
@@ -37,6 +39,9 @@ jobs:
       - name: Format check with flake8
         run: |
           flake8
+      - name: Type check with mypy
+        run: |
+          mypy src/dearprudence
       - name: Test with pytest
         run: |
           pytest -v --cov=./src/dearprudence --cov-report term-missing

CHANGELOG.md

@@ -5,3 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- Type annotations and mypy testing in CI.
+### Removed
+- `DtrRun` has been removed. Support for DTR files has been dropped.

pyproject.toml

@@ -10,3 +10,11 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools_scm]
 fallback_version = "999"
 write_to = "src/dearprudence/_version.py"
+
+[tool.mypy]
+python_version = "3.10"
+warn_unused_configs = true
+
+[[tool.mypy.overrides]]
+module = "intake"
+ignore_missing_imports = true

requirements-dev.txt

@@ -2,6 +2,7 @@ black
 build
 flake8
 intake_esm
+mypy
 pytest
 pytest-cov
 twine

src/dearprudence/core.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 
 
-__all__ = ["Cmip6Record", "SimpleRun", "DtrRun"]
+__all__ = ["Cmip6Record", "SimpleRun"]
 
 
 @dataclass
@@ -23,11 +23,3 @@ class SimpleRun:
     variable_id: str
     historical: Cmip6Record
     ssp: Cmip6Record
-
-
-@dataclass
-class DtrRun:
-    target: str
-    variable_id: str
-    tasmin: SimpleRun
-    tasmax: SimpleRun

src/dearprudence/io.py

@@ -1,13 +1,37 @@
 import dataclasses
 import json
+from os import PathLike
+from typing import Union, TextIO, BinaryIO, Any, TypedDict, Sequence
 
-from dearprudence.core import Cmip6Record, DtrRun, SimpleRun
+from dearprudence.core import Cmip6Record, SimpleRun
 
 
 __all__ = ["read_params", "write_params"]
 
 
-def _load_paramfile(urlpath):
+# TypedDict mappings representing intermediate dicts loaded from JSON.
+class Cmip6RecordMapping(TypedDict):
+    activity_id: str
+    experiment_id: str
+    table_id: str
+    variable_id: str
+    source_id: str
+    institution_id: str
+    member_id: str
+    grid_label: str
+    version: str
+
+
+class SimpleRunMapping(TypedDict):
+    variable_id: str
+    target: str
+    historical: Cmip6RecordMapping
+    ssp: Cmip6RecordMapping
+
+
+def _load_paramfile(
+    urlpath: Union[str, Union[TextIO, BinaryIO]]
+) -> Sequence[SimpleRunMapping]:
     # First readline() to pop-off and discard the first yaml bit of
     # the file, then load as JSON str. Keeps us from depending
     # on pyyaml.
@@ -22,7 +46,7 @@ def _load_paramfile(urlpath):
         return json.load(urlpath)
 
 
-def _unpack_simplerun(p):
+def _unpack_simplerun(p: SimpleRunMapping) -> SimpleRun:
     return SimpleRun(
         target=p["target"],
         variable_id=p["variable_id"],
@@ -31,37 +55,26 @@ def _unpack_simplerun(p):
     )
 
 
-def read_params(urlpath):
+def read_params(urlpath: Union[str, Union[TextIO, BinaryIO]]) -> list[SimpleRun]:
     """Read run parameters form yaml file"""
-    payload = _load_paramfile(urlpath)
-
-    out = []
-    for entry in payload:
-        if entry["variable_id"].lower() == "dtr":
-            out.append(
-                DtrRun(
-                    target=entry["target"],
-                    variable_id=entry["variable_id"],
-                    tasmin=_unpack_simplerun(entry["tasmin"]),
-                    tasmax=_unpack_simplerun(entry["tasmax"]),
-                )
-            )
-        else:
-            out.append(_unpack_simplerun(entry))
-
-    return out
+    return [_unpack_simplerun(x) for x in _load_paramfile(urlpath)]
 
 
 class _DataclassJSONEncoder(json.JSONEncoder):
     """Encoder to dump dataclasses to JSON"""
 
-    def default(self, o):
+    def default(self, o: Any) -> Any:
         if dataclasses.is_dataclass(o):
             return dataclasses.asdict(o)
         return super().default(o)
 
 
-def write_params(urlpath, runlist, mode="w", pretty=True):
+def write_params(
+    urlpath: Union[Union[str, bytes, PathLike[str], PathLike[bytes]], int],
+    runlist: Sequence[SimpleRun],
+    mode: str = "w",
+    pretty: bool = True,
+) -> None:
     """Write runs parameters to parameter file"""
     runlist = list(runlist)
 

src/dearprudence/utils.py

@@ -1,5 +1,6 @@
 from functools import cache
 
+from dearprudence.core import Cmip6Record
 from dearprudence.errors import (
     Cmip6CatalogNoEntriesError,
     Cmip6CatalogMultipleEntriesError,
@@ -11,7 +12,7 @@
 
 @cache
 def esm_datastore(
-    json_url="https://storage.googleapis.com/cmip6/pangeo-cmip6-noQC.json",
+    json_url: str = "https://storage.googleapis.com/cmip6/pangeo-cmip6-noQC.json",
 ):
     """
     Sugar to create an intake_esm.core.esm_datastore to pass into `in_cmip6_catalog`
@@ -31,7 +32,7 @@ def esm_datastore(
     return intake.open_esm_datastore(json_url)
 
 
-def cmip6_catalog_has(x, datastore=None):
+def cmip6_catalog_has(x: Cmip6Record, datastore=None) -> bool:
     """Check that Cmip6Record has an entry in CMIP6-In-The-Cloud catalog
 
     This requires the ``intake-esm`` package to be installed.