anemoi-utils 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

anemoi/utils/dates.py

@@ -11,11 +11,29 @@
 import calendar
 import datetime
 import re
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Set
+from typing import Tuple
+from typing import Union
 
 import aniso8601
 
 
-def normalise_frequency(frequency):
+def normalise_frequency(frequency: Union[int, str]) -> int:
+    """Normalise frequency to hours.
+
+    Parameters
+    ----------
+    frequency : int or str
+        The frequency to normalise.
+
+    Returns
+    -------
+    int
+        The normalised frequency in hours.
+    """
     if isinstance(frequency, int):
         return frequency
     assert isinstance(frequency, str), (type(frequency), frequency)
@@ -25,7 +43,7 @@ def normalise_frequency(frequency):
     return {"h": v, "d": v * 24}[unit]
 
 
-def _no_time_zone(date) -> datetime.datetime:
+def _no_time_zone(date: datetime.datetime) -> datetime.datetime:
     """Remove time zone information from a date.
 
     Parameters
@@ -43,7 +61,7 @@ def _no_time_zone(date) -> datetime.datetime:
 
 
 # this function is use in anemoi-datasets
-def as_datetime(date, keep_time_zone=False) -> datetime.datetime:
+def as_datetime(date: Union[datetime.date, datetime.datetime, str], keep_time_zone: bool = False) -> datetime.datetime:
     """Convert a date to a datetime object, removing any time zone information.
 
     Parameters
@@ -73,7 +91,23 @@ def as_datetime(date, keep_time_zone=False) -> datetime.datetime:
     raise ValueError(f"Invalid date type: {type(date)}")
 
 
-def _as_datetime_list(date, default_increment):
+def _as_datetime_list(
+    date: Union[datetime.date, datetime.datetime, str], default_increment: datetime.timedelta
+) -> iter:
+    """Convert a date to a list of datetime objects.
+
+    Parameters
+    ----------
+    date : datetime.date or datetime.datetime or str
+        The date to convert.
+    default_increment : datetime.timedelta
+        The default increment for the list.
+
+    Returns
+    -------
+    iter
+        An iterator of datetime objects.
+    """
     if isinstance(date, (list, tuple)):
         for d in date:
             yield from _as_datetime_list(d, default_increment)
@@ -102,12 +136,28 @@ def _as_datetime_list(date, default_increment):
     yield as_datetime(date)
 
 
-def as_datetime_list(date, default_increment=1):
+def as_datetime_list(
+    date: Union[datetime.date, datetime.datetime, str], default_increment: int = 1
+) -> list[datetime.datetime]:
+    """Convert a date to a list of datetime objects.
+
+    Parameters
+    ----------
+    date : datetime.date or datetime.datetime or str
+        The date to convert.
+    default_increment : int, optional
+        The default increment in hours, by default 1.
+
+    Returns
+    -------
+    list of datetime.datetime
+        A list of datetime objects.
+    """
     default_increment = frequency_to_timedelta(default_increment)
     return list(_as_datetime_list(date, default_increment))
 
 
-def as_timedelta(frequency) -> datetime.timedelta:
+def as_timedelta(frequency: Union[int, str, datetime.timedelta]) -> datetime.timedelta:
     """Convert anything to a timedelta object.
 
     Parameters
@@ -155,12 +205,19 @@ def as_timedelta(frequency) -> datetime.timedelta:
         unit = {"h": "hours", "d": "days", "s": "seconds", "m": "minutes"}[unit]
         return datetime.timedelta(**{unit: v})
 
-    m = frequency.split(":")
-    if len(m) == 2:
-        return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]))
+    if re.match(r"^\d+:\d+(:\d+)?$", frequency):
+        m = frequency.split(":")
+        if len(m) == 2:
+            return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]))
+
+        if len(m) == 3:
+            return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]), seconds=int(m[2]))
 
-    if len(m) == 3:
-        return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]), seconds=int(m[2]))
+    if re.match(r"^\d+ days?, \d+:\d+:\d+$", frequency):
+        m = frequency.split(", ")
+        days = int(m[0].split()[0])
+        hms = m[1].split(":")
+        return datetime.timedelta(days=days, hours=int(hms[0]), minutes=int(hms[1]), seconds=int(hms[2]))
 
     # ISO8601
     try:
@@ -171,12 +228,23 @@ def as_timedelta(frequency) -> datetime.timedelta:
     raise ValueError(f"Cannot convert frequency {frequency} to timedelta")
 
 
-def frequency_to_timedelta(frequency) -> datetime.timedelta:
-    """Convert a frequency to a timedelta object."""
+def frequency_to_timedelta(frequency: Union[int, str, datetime.timedelta]) -> datetime.timedelta:
+    """Convert a frequency to a timedelta object.
+
+    Parameters
+    ----------
+    frequency : int or str or datetime.timedelta
+        The frequency to convert.
+
+    Returns
+    -------
+    datetime.timedelta
+        The timedelta object.
+    """
     return as_timedelta(frequency)
 
 
-def frequency_to_string(frequency) -> str:
+def frequency_to_string(frequency: datetime.timedelta) -> str:
     """Convert a frequency (i.e. a datetime.timedelta) to a string.
 
     Parameters
@@ -223,20 +291,19 @@ def frequency_to_string(frequency) -> str:
     return str(frequency)
 
 
-def frequency_to_seconds(frequency) -> int:
+def frequency_to_seconds(frequency: Union[int, str, datetime.timedelta]) -> int:
     """Convert a frequency to seconds.
 
     Parameters
     ----------
-    frequency : _type_
-        _description_
+    frequency : int or str or datetime.timedelta
+        The frequency to convert.
 
     Returns
     -------
     int
         Number of seconds.
     """
-
     result = frequency_to_timedelta(frequency).total_seconds()
     assert int(result) == result, result
     return int(result)
@@ -269,7 +336,19 @@ MONTH = {
 }
 
 
-def _make_day(day):
+def _make_day(day: Optional[Tuple[int, List[int]]]) -> Set[int]:
+    """Create a set of days.
+
+    Parameters
+    ----------
+    day : int or list of int or None
+        The day(s) to include in the set.
+
+    Returns
+    -------
+    set of int
+        A set of days.
+    """
     if day is None:
         return set(range(1, 32))
     if not isinstance(day, list):
@@ -277,7 +356,19 @@ def _make_day(day):
     return {int(d) for d in day}
 
 
-def _make_week(week):
+def _make_week(week: Optional[Tuple[str, List[str]]]) -> Set[int]:
+    """Create a set of weekdays.
+
+    Parameters
+    ----------
+    week : str or list of str or None
+        The weekday(s) to include in the set.
+
+    Returns
+    -------
+    set of int
+        A set of weekdays.
+    """
     if week is None:
         return set(range(7))
     if not isinstance(week, list):
@@ -285,7 +376,19 @@ def _make_week(week):
     return {DOW[w.lower()] for w in week}
 
 
-def _make_months(months):
+def _make_months(months: Optional[Union[int, str, List[Union[int, str]]]]) -> Set[int]:
+    """Create a set of months.
+
+    Parameters
+    ----------
+    months : int or str or list of int or str or None
+        The month(s) to include in the set.
+
+    Returns
+    -------
+    set of int
+        A set of months.
+    """
     if months is None:
         return set(range(1, 13))
 
@@ -298,23 +401,32 @@ def _make_months(months):
 class DateTimes:
     """The DateTimes class is an iterator that generates datetime objects within a given range."""
 
-    def __init__(self, start, end, increment=24, *, day_of_month=None, day_of_week=None, calendar_months=None):
-        """_summary_
+    def __init__(
+        self,
+        start: Union[datetime.date, datetime.datetime, str],
+        end: Union[datetime.date, datetime.datetime, str],
+        increment: int = 24,
+        *,
+        day_of_month: Optional[Tuple[int, List[int]]] = None,
+        day_of_week: Optional[Tuple[str, List[str]]] = None,
+        calendar_months: Optional[Union[int, str, List[Union[int, str]]]] = None,
+    ):
+        """Initialize the DateTimes iterator.
 
         Parameters
         ----------
-        start : _type_
-            _description_
-        end : _type_
-            _description_
+        start : datetime.date or datetime.datetime or str
+            The start date.
+        end : datetime.date or datetime.datetime or str
+            The end date.
         increment : int, optional
-            _description_, by default 24
-        day_of_month : _type_, optional
-            _description_, by default None
-        day_of_week : _type_, optional
-            _description_, by default None
-        calendar_months : _type_, optional
-            _description_, by default None
+            The increment in hours, by default 24.
+        day_of_month : int or list of int or None, optional
+            The day(s) of the month to include, by default None.
+        day_of_week : str or list of str or None, optional
+            The day(s) of the week to include, by default None.
+        calendar_months : int or str or list of int or str or None, optional
+            The month(s) to include, by default None.
         """
         self.start = as_datetime(start)
         self.end = as_datetime(end)
@@ -323,7 +435,14 @@ class DateTimes:
         self.day_of_week = _make_week(day_of_week)
         self.calendar_months = _make_months(calendar_months)
 
-    def __iter__(self):
+    def __iter__(self) -> iter:
+        """Iterate over the datetime objects.
+
+        Returns
+        -------
+        iter
+            An iterator of datetime objects.
+        """
         date = self.start
         while date <= self.end:
             if (
@@ -339,13 +458,13 @@ class DateTimes:
 class Year(DateTimes):
     """Year is defined as the months of January to December."""
 
-    def __init__(self, year, **kwargs):
-        """_summary_
+    def __init__(self, year: int, **kwargs):
+        """Initialize the Year iterator.
 
         Parameters
         ----------
         year : int
-            _description_
+            The year.
         """
         super().__init__(datetime.datetime(year, 1, 1), datetime.datetime(year, 12, 31), **kwargs)
 
@@ -353,13 +472,13 @@ class Year(DateTimes):
 class Winter(DateTimes):
     """Winter is defined as the months of December, January and February."""
 
-    def __init__(self, year, **kwargs):
-        """_summary_
+    def __init__(self, year: int, **kwargs):
+        """Initialize the Winter iterator.
 
         Parameters
         ----------
         year : int
-            _description_
+            The year.
         """
         super().__init__(
             datetime.datetime(year, 12, 1),
@@ -371,13 +490,13 @@ class Winter(DateTimes):
 class Spring(DateTimes):
     """Spring is defined as the months of March, April and May."""
 
-    def __init__(self, year, **kwargs):
-        """_summary_
+    def __init__(self, year: int, **kwargs):
+        """Initialize the Spring iterator.
 
         Parameters
         ----------
         year : int
-            _description_
+            The year.
         """
         super().__init__(datetime.datetime(year, 3, 1), datetime.datetime(year, 5, 31), **kwargs)
 
@@ -385,13 +504,13 @@ class Spring(DateTimes):
 class Summer(DateTimes):
     """Summer is defined as the months of June, July and August."""
 
-    def __init__(self, year, **kwargs):
-        """_summary_
+    def __init__(self, year: int, **kwargs):
+        """Initialize the Summer iterator.
 
         Parameters
         ----------
         year : int
-            _description_
+            The year.
         """
         super().__init__(datetime.datetime(year, 6, 1), datetime.datetime(year, 8, 31), **kwargs)
 
@@ -399,13 +518,13 @@ class Summer(DateTimes):
 class Autumn(DateTimes):
     """Autumn is defined as the months of September, October and November."""
 
-    def __init__(self, year, **kwargs):
-        """_summary_
+    def __init__(self, year: int, **kwargs):
+        """Initialize the Autumn iterator.
 
         Parameters
         ----------
         year : int
-            _description_
+            The year.
         """
         super().__init__(datetime.datetime(year, 9, 1), datetime.datetime(year, 11, 30), **kwargs)
 
@@ -413,13 +532,27 @@ class Autumn(DateTimes):
 class ConcatDateTimes:
     """ConcatDateTimes is an iterator that generates datetime objects from a list of dates."""
 
-    def __init__(self, *dates):
+    def __init__(self, *dates: DateTimes):
+        """Initialize the ConcatDateTimes iterator.
+
+        Parameters
+        ----------
+        dates : DateTimes
+            The list of DateTimes objects.
+        """
         if len(dates) == 1 and isinstance(dates[0], list):
             dates = dates[0]
 
         self.dates = dates
 
-    def __iter__(self):
+    def __iter__(self) -> iter:
+        """Iterate over the datetime objects.
+
+        Returns
+        -------
+        iter
+            An iterator of datetime objects.
+        """
         for date in self.dates:
             yield from date
 
@@ -427,15 +560,43 @@ class ConcatDateTimes:
 class EnumDateTimes:
     """EnumDateTimes is an iterator that generates datetime objects from a list of dates."""
 
-    def __init__(self, dates):
+    def __init__(self, dates: list[Union[datetime.date, datetime.datetime, str]]):
+        """Initialize the EnumDateTimes iterator.
+
+        Parameters
+        ----------
+        dates : list of datetime.date or datetime.datetime or str
+            The list of dates.
+        """
         self.dates = dates
 
-    def __iter__(self):
+    def __iter__(self) -> iter:
+        """Iterate over the datetime objects.
+
+        Returns
+        -------
+        iter
+            An iterator of datetime objects.
+        """
         for date in self.dates:
             yield as_datetime(date)
 
 
-def datetimes_factory(*args, **kwargs):
+def datetimes_factory(*args: Any, **kwargs: Any) -> Union[DateTimes, ConcatDateTimes, EnumDateTimes]:
+    """Create a DateTimes, ConcatDateTimes, or EnumDateTimes object.
+
+    Parameters
+    ----------
+    *args : Any
+        Positional arguments.
+    **kwargs : Any
+        Keyword arguments.
+
+    Returns
+    -------
+    DateTimes or ConcatDateTimes or EnumDateTimes
+        The created object.
+    """
     if args and kwargs:
         raise ValueError("Cannot provide both args and kwargs for a list of dates")
 

anemoi/utils/devtools.py

@@ -8,29 +8,74 @@
 # nor does it submit to any jurisdiction.
 
 
+from typing import Any
+
 import cartopy.crs as ccrs
 import cartopy.feature as cfeature
 import matplotlib.pyplot as plt
 import matplotlib.tri as tri
 import numpy as np
 
-"""FOR DEVELOPMENT PURPOSES ONLY
+"""FOR DEVELOPMENT PURPOSES ONLY.
 
 This module contains
-
 """
 
 # TODO: use earthkit-plots
 
 
-def fix(lons):
+def fix(lons: np.ndarray) -> np.ndarray:
+    """Fix longitudes greater than 180 degrees.
+
+    Parameters
+    ----------
+    lons : np.ndarray
+        Array of longitudes.
+
+    Returns
+    -------
+    np.ndarray
+        Array of fixed longitudes.
+    """
     return np.where(lons > 180, lons - 360, lons)
 
 
 def plot_values(
-    values, latitudes, longitudes, title=None, missing_value=None, min_value=None, max_value=None, **kwargs
-):
-
+    values: np.ndarray,
+    latitudes: np.ndarray,
+    longitudes: np.ndarray,
+    title: str = None,
+    missing_value: float = None,
+    min_value: float = None,
+    max_value: float = None,
+    **kwargs: dict,
+) -> plt.Axes:
+    """Plot values on a map.
+
+    Parameters
+    ----------
+    values : np.ndarray
+        Array of values to plot.
+    latitudes : np.ndarray
+        Array of latitudes.
+    longitudes : np.ndarray
+        Array of longitudes.
+    title : str, optional
+        Title of the plot, by default None.
+    missing_value : float, optional
+        Value to use for missing data, by default None.
+    min_value : float, optional
+        Minimum value for the plot, by default None.
+    max_value : float, optional
+        Maximum value for the plot, by default None.
+    **kwargs : dict
+        Additional keyword arguments for the plot.
+
+    Returns
+    -------
+    plt.Axes
+        The plot axes.
+    """
     _, ax = plt.subplots(subplot_kw={"projection": ccrs.PlateCarree()})
     ax.coastlines()
     ax.add_feature(cfeature.BORDERS, linestyle=":")
@@ -77,7 +122,23 @@ def plot_values(
     return ax
 
 
-def plot_field(field, title=None, **kwargs):
+def plot_field(field: Any, title: str = None, **kwargs: dict) -> plt.Axes:
+    """Plot a field on a map.
+
+    Parameters
+    ----------
+    field : Any
+        The field to plot.
+    title : str, optional
+        Title of the plot, by default None.
+    **kwargs : dict
+        Additional keyword arguments for the plot.
+
+    Returns
+    -------
+    plt.Axes
+        The plot axes.
+    """
     values = field.to_numpy(flatten=True)
     latitudes, longitudes = field.grid_points()
     return plot_values(values, latitudes, longitudes, title=title, **kwargs)

anemoi/utils/grib.py

@@ -11,11 +11,12 @@
 """Utilities for working with GRIB parameters.
 
 See https://codes.ecmwf.int/grib/param-db/ for more information.
-
 """
 
 import logging
 import re
+from typing import Dict
+from typing import Union
 
 import requests
 
@@ -25,7 +26,14 @@ LOG = logging.getLogger(__name__)
 
 
 @cached(collection="grib", expires=30 * 24 * 60 * 60)
-def _units():
+def _units() -> Dict[str, str]:
+    """Fetch and cache GRIB parameter units.
+
+    Returns
+    -------
+    dict
+        A dictionary mapping unit ids to their names.
+    """
     r = requests.get("https://codes.ecmwf.int/parameter-database/api/v1/unit/")
     r.raise_for_status()
     units = r.json()
@@ -33,7 +41,24 @@ def _units():
 
 
 @cached(collection="grib", expires=30 * 24 * 60 * 60)
-def _search_param(name):
+def _search_param(name: str) -> Dict[str, Union[str, int]]:
+    """Search for a GRIB parameter by name.
+
+    Parameters
+    ----------
+    name : str
+        Parameter name to search for.
+
+    Returns
+    -------
+    dict
+        A dictionary containing parameter details.
+
+    Raises
+    ------
+    KeyError
+        If no parameter is found.
+    """
     name = re.escape(name)
     r = requests.get(f"https://codes.ecmwf.int/parameter-database/api/v1/param/?search=^{name}$&regex=true")
     r.raise_for_status()
@@ -68,7 +93,6 @@ def shortname_to_paramid(shortname: str) -> int:
 
     >>> shortname_to_paramid("2t")
     167
-
     """
     return _search_param(shortname)["id"]
 
@@ -88,12 +112,11 @@ def paramid_to_shortname(paramid: int) -> str:
 
     >>> paramid_to_shortname(167)
     '2t'
-
     """
     return _search_param(str(paramid))["shortname"]
 
 
-def units(param) -> str:
+def units(param: Union[int, str]) -> str:
     """Return the units of a GRIB parameter given its name or id.
 
     Parameters
@@ -108,14 +131,13 @@ def units(param) -> str:
 
     >>> unit(167)
     'K'
-
     """
 
     unit_id = str(_search_param(str(param))["unit_id"])
     return _units()[unit_id]
 
 
-def must_be_positive(param) -> bool:
+def must_be_positive(param: Union[int, str]) -> bool:
     """Check if a parameter must be positive.
 
     Parameters
@@ -130,6 +152,5 @@ def must_be_positive(param) -> bool:
 
     >>> must_be_positive("tp")
     True
-
     """
     return units(param) in ["m", "kg kg**-1", "m of water equivalent"]