From d49ba5d25e9bc22757500703b4be1a60b5fc0c72 Mon Sep 17 00:00:00 2001 From: arnavsharma990 <2006arnavsharma@gmail.com> Date: Fri, 20 Feb 2026 12:53:05 +0530 Subject: [PATCH 1/3] style: add type hints to public functions in ops and config - loading.py: load_input_dataset(fp: str) -> xr.Dataset - selection.py: select_by_kwargs, check_point_in_dataset, check_step; _normalize_slice_startstop, _normalize_slice_step - subsetting.py: extract_variable with Union/Optional types; use None default for coords_to_sample to support Optional type - mapping.py: map_dims_and_variables, _check_for_malformed_list_arg; dim_mapping: Dict[str, DimMapping] - stacking.py: stack_variables_as_coord_values, stack_variables_by_coord_values - chunking.py: check_chunk_size, chunk_dataset with Dict[str, int] - config.py: validate_config(config_inputs: Dict[str, InputDataset]) -> None Co-authored-by: Cursor --- mllam_data_prep/config.py | 2 +- mllam_data_prep/ops/chunking.py | 7 +++++-- mllam_data_prep/ops/loading.py | 2 +- mllam_data_prep/ops/mapping.py | 13 +++++++++++-- mllam_data_prep/ops/selection.py | 22 +++++++++++++++++----- mllam_data_prep/ops/stacking.py | 11 +++++++++-- mllam_data_prep/ops/subsetting.py | 13 ++++++++++++- 7 files changed, 56 insertions(+), 14 deletions(-) diff --git a/mllam_data_prep/config.py b/mllam_data_prep/config.py index b7a0ab4..72c774f 100644 --- a/mllam_data_prep/config.py +++ b/mllam_data_prep/config.py @@ -12,7 +12,7 @@ class InvalidConfigException(Exception): pass -def validate_config(config_inputs): +def validate_config(config_inputs: Dict[str, "InputDataset"]) -> None: """ Validate that, in the config: - either `variables` or `derived_variables` are present in the config diff --git a/mllam_data_prep/ops/chunking.py b/mllam_data_prep/ops/chunking.py index 9df27e9..cc61312 100644 --- a/mllam_data_prep/ops/chunking.py +++ b/mllam_data_prep/ops/chunking.py @@ -1,11 +1,14 @@ +from typing import Dict + import numpy as np +import xarray as xr from loguru import logger # Max chunk size warning CHUNK_MAX_SIZE_WARNING = 1 * 1024**3 # 1GB -def check_chunk_size(ds, chunks): +def check_chunk_size(ds: xr.Dataset, chunks: Dict[str, int]) -> None: """ Check the chunk size and warn if it exceeds CHUNK_MAX_SIZE_WARNING. @@ -44,7 +47,7 @@ def check_chunk_size(ds, chunks): ) -def chunk_dataset(ds, chunks): +def chunk_dataset(ds: xr.Dataset, chunks: Dict[str, int]) -> xr.Dataset: """ Check the chunk size and chunk the dataset. diff --git a/mllam_data_prep/ops/loading.py b/mllam_data_prep/ops/loading.py index f6bfc34..b104c5d 100644 --- a/mllam_data_prep/ops/loading.py +++ b/mllam_data_prep/ops/loading.py @@ -1,7 +1,7 @@ import xarray as xr -def load_input_dataset(fp): +def load_input_dataset(fp: str) -> xr.Dataset: """ Load the dataset diff --git a/mllam_data_prep/ops/mapping.py b/mllam_data_prep/ops/mapping.py index 9482ff8..79fe57c 100644 --- a/mllam_data_prep/ops/mapping.py +++ b/mllam_data_prep/ops/mapping.py @@ -1,14 +1,23 @@ +from typing import Any, Dict, List + +import xarray as xr + +from ..config import DimMapping from .stacking import stack_variables_as_coord_values, stack_variables_by_coord_values -def _check_for_malformed_list_arg(s): +def _check_for_malformed_list_arg(s: Any) -> None: if isinstance(s, str) and "," in s: raise Exception( "Rather than writing `{s}` to define a list you would `[{s}]` in the config file." ) -def map_dims_and_variables(ds, dim_mapping, expected_input_var_dims): +def map_dims_and_variables( + ds: xr.Dataset, + dim_mapping: Dict[str, DimMapping], + expected_input_var_dims: List[str], +) -> xr.DataArray: """ Map the input dimensions to the architecture dimensions using the `dim_mapping` dictionary. Each key in the `dim_mapping` diff --git a/mllam_data_prep/ops/selection.py b/mllam_data_prep/ops/selection.py index 37b91c1..d2b9a22 100644 --- a/mllam_data_prep/ops/selection.py +++ b/mllam_data_prep/ops/selection.py @@ -1,11 +1,15 @@ import datetime +from typing import Any, List, Union import pandas as pd +import xarray as xr from ..config import Range -def _normalize_slice_startstop(s): +def _normalize_slice_startstop( + s: Union[pd.Timestamp, str, Any] +) -> Union[pd.Timestamp, str, Any]: if isinstance(s, pd.Timestamp): return s elif isinstance(s, str): @@ -17,7 +21,9 @@ def _normalize_slice_startstop(s): return s -def _normalize_slice_step(s): +def _normalize_slice_step( + s: Union[pd.Timedelta, str, Any] +) -> Union[pd.Timedelta, str, Any]: if isinstance(s, pd.Timedelta): return s elif isinstance(s, str): @@ -29,7 +35,9 @@ def _normalize_slice_step(s): return s -def select_by_kwargs(ds, **coord_ranges): +def select_by_kwargs( + ds: xr.Dataset, **coord_ranges: Union[Range, List[Any]] +) -> xr.Dataset: """ Do `xr.Dataset.sel` on `ds` using the `coord_ranges` to select the coordinates, for each entry in the dictionary, the key is the coordinate name and the value is the selection @@ -92,7 +100,7 @@ def select_by_kwargs(ds, **coord_ranges): return ds -def check_point_in_dataset(coord, point, ds): +def check_point_in_dataset(coord: str, point: Any, ds: xr.Dataset) -> None: """ check that the requested point is in the data. """ @@ -102,7 +110,11 @@ def check_point_in_dataset(coord, point, ds): ) -def check_step(sel_step, coord, ds): +def check_step( + sel_step: Union[pd.Timedelta, datetime.timedelta], + coord: str, + ds: xr.Dataset, +) -> None: """ check that the step requested is exactly what the data has """ diff --git a/mllam_data_prep/ops/stacking.py b/mllam_data_prep/ops/stacking.py index a56e0fd..4a436ef 100644 --- a/mllam_data_prep/ops/stacking.py +++ b/mllam_data_prep/ops/stacking.py @@ -1,7 +1,9 @@ import xarray as xr -def stack_variables_as_coord_values(ds, name_format, combined_dim_name): +def stack_variables_as_coord_values( + ds: xr.Dataset, name_format: str, combined_dim_name: str +) -> xr.DataArray: """ combine all variables in an xr.Dataset into a single xr.DataArray by stacking the variables along a new coordinate with the name given @@ -50,7 +52,12 @@ def stack_variables_as_coord_values(ds, name_format, combined_dim_name): return da_combined -def stack_variables_by_coord_values(ds, coord, name_format, combined_dim_name): +def stack_variables_by_coord_values( + ds: xr.Dataset, + coord: str, + name_format: str, + combined_dim_name: str, +) -> xr.DataArray: """ combine all variables in an xr.Dataset on all coordinate values of `coord` into a single xr.DataArray diff --git a/mllam_data_prep/ops/subsetting.py b/mllam_data_prep/ops/subsetting.py index 80f2ce1..abc7702 100644 --- a/mllam_data_prep/ops/subsetting.py +++ b/mllam_data_prep/ops/subsetting.py @@ -1,4 +1,13 @@ -def extract_variable(ds, var_name, coords_to_sample=dict()): +from typing import Any, Dict, List, Optional, Union + +import xarray as xr + + +def extract_variable( + ds: xr.Dataset, + var_name: Union[str, Dict[str, Any], List[str]], + coords_to_sample: Optional[Dict[str, Any]] = None, +) -> xr.DataArray: """ Extract specified variable from the provided input dataset. If coordinates for subsetting are defined, then subset the variable along @@ -21,6 +30,8 @@ def extract_variable(ds, var_name, coords_to_sample=dict()): da: xr.DataArray Extracted variable (subsetted along the specified coordinates) """ + if coords_to_sample is None: + coords_to_sample = {} try: da = ds[var_name] From 149031bb307cdfcf6b76cc07dfaa591b2878e321 Mon Sep 17 00:00:00 2001 From: arnavsharma990 <2006arnavsharma@gmail.com> Date: Fri, 20 Feb 2026 13:20:24 +0530 Subject: [PATCH 2/3] fix: correct type hints for extract_variable based on actual usage - Change var_name type from Union[str, Dict, List] to str (var_name is always a string when used to index ds[var_name]) - Change coords_to_sample type from Dict[str, Any] to Dict[str, ValueSelection] (matches actual usage where sampling.values and sampling.units are accessed) - Update docstring to reflect correct types - Import ValueSelection from ..config Co-authored-by: Cursor --- mllam_data_prep/ops/subsetting.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/mllam_data_prep/ops/subsetting.py b/mllam_data_prep/ops/subsetting.py index abc7702..7c7cdcd 100644 --- a/mllam_data_prep/ops/subsetting.py +++ b/mllam_data_prep/ops/subsetting.py @@ -1,12 +1,14 @@ -from typing import Any, Dict, List, Optional, Union +from typing import Dict, Optional import xarray as xr +from ..config import ValueSelection + def extract_variable( ds: xr.Dataset, - var_name: Union[str, Dict[str, Any], List[str]], - coords_to_sample: Optional[Dict[str, Any]] = None, + var_name: str, + coords_to_sample: Optional[Dict[str, ValueSelection]] = None, ) -> xr.DataArray: """ Extract specified variable from the provided input dataset. If @@ -17,13 +19,12 @@ def extract_variable( ---------- ds : xr.Dataset Input dataset - var_name : Union[Dict, List] - Either a list or dictionary with variables to extract. - If a dictionary the keys are the variable name and the values are - entries for each coordinate and coordinate values to extract - coords_to_sample: Dict + var_name : str + Name of the variable to extract from the dataset + coords_to_sample: Dict[str, ValueSelection], optional Optional argument for subsetting/sampling along the specified - coordinates + coordinates. Keys are coordinate names, values are ValueSelection + objects defining the values to select and optionally the units. Returns ---------- From 75cd6b75fd64508a18dbca5c3f0d08aa973ca84a Mon Sep 17 00:00:00 2001 From: arnavsharma990 <2006arnavsharma@gmail.com> Date: Fri, 20 Feb 2026 13:22:48 +0530 Subject: [PATCH 3/3] fix: improve type hints for _normalize_slice functions - _normalize_slice_startstop: Change from Union[pd.Timestamp, str, Any] to Union[pd.Timestamp, str, int, float] to match Range.start/end types - _normalize_slice_step: Change from Union[pd.Timedelta, str, Any] to Optional[Union[pd.Timedelta, str, int, float]] to match Range.step (which is Optional) and account for None values - Remove unused Any import, add Optional import Co-authored-by: Cursor --- mllam_data_prep/ops/selection.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/mllam_data_prep/ops/selection.py b/mllam_data_prep/ops/selection.py index d2b9a22..e73530f 100644 --- a/mllam_data_prep/ops/selection.py +++ b/mllam_data_prep/ops/selection.py @@ -1,5 +1,5 @@ import datetime -from typing import Any, List, Union +from typing import List, Optional, Union import pandas as pd import xarray as xr @@ -8,8 +8,8 @@ def _normalize_slice_startstop( - s: Union[pd.Timestamp, str, Any] -) -> Union[pd.Timestamp, str, Any]: + s: Union[pd.Timestamp, str, int, float] +) -> Union[pd.Timestamp, str, int, float]: if isinstance(s, pd.Timestamp): return s elif isinstance(s, str): @@ -22,8 +22,8 @@ def _normalize_slice_startstop( def _normalize_slice_step( - s: Union[pd.Timedelta, str, Any] -) -> Union[pd.Timedelta, str, Any]: + s: Optional[Union[pd.Timedelta, str, int, float]] +) -> Optional[Union[pd.Timedelta, str, int, float]]: if isinstance(s, pd.Timedelta): return s elif isinstance(s, str):