Add types

Author: spencerkclark

xarray/coding/cftime_offsets.py

@@ -45,7 +45,7 @@
 from datetime import datetime, timedelta
 from enum import Enum
 from functools import partial
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 import numpy as np
 import pandas as pd
@@ -67,6 +67,10 @@
     cftime = None
 
 
+if TYPE_CHECKING:
+    from xarray.core.types import InclusiveOptions, SideOptions
+
+
 def get_date_type(calendar, use_cftime=True):
     """Return the cftime date type for a given calendar name."""
     if cftime is None:
@@ -903,8 +907,8 @@ def cftime_range(
     freq="D",
     normalize=False,
     name=None,
-    closed=_NoDefault,
-    inclusive=None,
+    closed: _NoDefault | SideOptions = _NoDefault,
+    inclusive: None | InclusiveOptions = None,
     calendar="standard",
 ):
     """Return a fixed frequency CFTimeIndex.
@@ -945,23 +949,22 @@ def cftime_range(
     Notes
     -----
     This function is an analog of ``pandas.date_range`` for use in generating
-    sequences of ``cftime.datetime`` objects.  It supports most of the features
-    of ``pandas.date_range`` (e.g. specifying how the index is ``closed`` on
-    either side, or whether or not to ``normalize`` the start and end bounds);
-    however, there are some notable exceptions:
-
+    sequences of ``cftime.datetime`` objects.  It supports most of the
+    features of ``pandas.date_range`` (e.g. specifying how the index is
+    ``closed`` on either side, or whether or not to ``normalize`` the start and
+    end bounds); however, there are some notable exceptions:
     - You cannot specify a ``tz`` (time zone) argument.
     - Start or end dates specified as partial-datetime strings must use the
       `ISO-8601 format <https://en.wikipedia.org/wiki/ISO_8601>`_.
     - It supports many, but not all, frequencies supported by
       ``pandas.date_range``.  For example it does not currently support any of
       the business-related or semi-monthly frequencies.
     - Compound sub-monthly frequencies are not supported, e.g. '1H1min', as
-      these can easily be written in terms of the finest common resolution, e.g.
-      '61min'.
+      these can easily be written in terms of the finest common resolution,
+      e.g. '61min'.
 
-    Valid simple frequency strings for use with ``cftime``-calendars include any
-    multiples of the following.
+    Valid simple frequency strings for use with ``cftime``-calendars include
+    any multiples of the following.
 
     +--------+--------------------------+
     | Alias  | Description              |
@@ -1140,8 +1143,8 @@ def date_range(
     tz=None,
     normalize=False,
     name=None,
-    closed=_NoDefault,
-    inclusive=None,
+    closed: _NoDefault | SideOptions = _NoDefault,
+    inclusive: None | InclusiveOptions = None,
     calendar="standard",
     use_cftime=None,
 ):

xarray/core/types.py

@@ -172,6 +172,7 @@ def dtype(self) -> np.dtype:
 
 CoarsenBoundaryOptions = Literal["exact", "trim", "pad"]
 SideOptions = Literal["left", "right"]
+InclusiveOptions = Literal["both", "neither", "left", "right"]
 
 ScaleOptions = Literal["linear", "symlog", "log", "logit", None]
 HueStyleOptions = Literal["continuous", "discrete", None]

xarray/tests/test_cftime_offsets.py

@@ -1411,7 +1411,7 @@ def as_timedelta_not_implemented_error():
 
 
 @pytest.mark.parametrize("function", [cftime_range, date_range])
-def test_cftime_or_date_range_closed_and_inclusive_error(function):
+def test_cftime_or_date_range_closed_and_inclusive_error(function) -> None:
     if function == cftime_range and not has_cftime:
         pytest.skip("requires cftime")
 
@@ -1420,7 +1420,7 @@ def test_cftime_or_date_range_closed_and_inclusive_error(function):
 
 
 @pytest.mark.parametrize("function", [cftime_range, date_range])
-def test_cftime_or_date_range_invalid_closed_value(function):
+def test_cftime_or_date_range_invalid_closed_value(function) -> None:
     if function == cftime_range and not has_cftime:
         pytest.skip("requires cftime")
 
@@ -1432,7 +1432,7 @@ def test_cftime_or_date_range_invalid_closed_value(function):
 @pytest.mark.parametrize(
     ("closed", "inclusive"), [(None, "both"), ("left", "left"), ("right", "right")]
 )
-def test_cftime_or_date_range_closed(function, closed, inclusive):
+def test_cftime_or_date_range_closed(function, closed, inclusive) -> None:
     if function == cftime_range and not has_cftime:
         pytest.skip("requires cftime")
 
@@ -1445,7 +1445,7 @@ def test_cftime_or_date_range_closed(function, closed, inclusive):
 
 
 @pytest.mark.parametrize("function", [cftime_range, date_range])
-def test_cftime_or_date_range_inclusive_None(function):
+def test_cftime_or_date_range_inclusive_None(function) -> None:
     if function == cftime_range and not has_cftime:
         pytest.skip("requires cftime")