enable mypy and fix types accordingly

.pre-commit-config.yaml

@@ -19,3 +19,10 @@ repos:
     rev: 5.0.4
     hooks:
       - id: flake8
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.14.1
+    hooks:
+      - id: mypy
+        additional_dependencies: [types-PyYAML, types-Pillow, types-tqdm]
+        description: Check for type errors
+        files: ^mllam_data_prep/

mllam_data_prep/__main__.py

@@ -11,7 +11,7 @@
     import psutil
     from dask.diagnostics import ProgressBar
     from dask.distributed import LocalCluster
-except ImportError or ModuleNotFoundError:
+except (ImportError, ModuleNotFoundError):
     DASK_DISTRIBUTED_AVAILABLE = False
 
 if __name__ == "__main__":

mllam_data_prep/config.py

@@ -235,7 +235,7 @@ class InputDataset:
     variables: Optional[Union[List[str], Dict[str, Dict[str, ValueSelection]]]] = None
     derived_variables: Optional[Dict[str, DerivedVariable]] = None
     attributes: Optional[Dict[str, Any]] = field(default_factory=dict)
-    coord_ranges: Dict[str, Range] = None
+    coord_ranges: Optional[Dict[str, Range]] = None
 
 
 @dataclass

mllam_data_prep/create_dataset.py

@@ -2,6 +2,7 @@
 import shutil
 from collections import defaultdict
 from pathlib import Path
+from typing import Optional
 
 import numpy as np
 import xarray as xr
@@ -294,7 +295,7 @@ def create_dataset(config: Config):
     return ds
 
 
-def create_dataset_zarr(fp_config, fp_zarr: str = None):
+def create_dataset_zarr(fp_config: Path, fp_zarr: Optional[str | Path] = None):
     """
     Create a dataset from the input datasets specified in the config file and write it to a zarr file.
     The path to the zarr file is the same as the config file, but with the extension changed to '.zarr'.

tests/test_config.py

@@ -119,3 +119,13 @@ def test_get_config_nested():
         assert input_config.target_output_variable is not None
         with pytest.raises(AttributeError):
             input_config.foobarfield
+
+
+def test_config_roundtrip():
+    original_config = mdp.Config.from_yaml(VALID_EXAMPLE_CONFIG_YAML)
+    roundtrip_config_dict = mdp.Config.from_dict(original_config.to_dict())
+    roundtrip_config_yaml = mdp.Config.from_yaml(original_config.to_yaml())
+    roundtrip_config_json = mdp.Config.from_json(original_config.to_json())
+    assert original_config == roundtrip_config_dict
+    assert original_config == roundtrip_config_yaml
+    assert original_config == roundtrip_config_json