ENH: Session type, parse_session, new ExchangeCalendar methods

exchange_calendars/calendar_helpers.py

@@ -1,12 +1,23 @@
+from __future__ import annotations
+import typing
+import datetime
+
 import numpy as np
 import pandas as pd
 
+from exchange_calendars import errors
+
+if typing.TYPE_CHECKING:
+    from exchange_calendars import ExchangeCalendar
+
 NANOSECONDS_PER_MINUTE = int(6e10)
 
 NP_NAT = np.array([pd.NaT], dtype=np.int64)[0]
 
+Session = typing.Union[pd.Timestamp, str, int, float, datetime.datetime]
 
-def next_divider_idx(dividers, minute_val):
+
+def next_divider_idx(dividers: np.ndarray, minute_val: int) -> int:
 
     divider_idx = np.searchsorted(dividers, minute_val, side="right")
     target = dividers[divider_idx]
@@ -18,7 +29,7 @@ def next_divider_idx(dividers, minute_val):
         return divider_idx
 
 
-def previous_divider_idx(dividers, minute_val):
+def previous_divider_idx(dividers: np.ndarray, minute_val: int) -> int:
 
     divider_idx = np.searchsorted(dividers, minute_val)
 
@@ -29,11 +40,11 @@ def previous_divider_idx(dividers, minute_val):
 
 
 def compute_all_minutes(
-    opens_in_ns,
-    break_starts_in_ns,
-    break_ends_in_ns,
-    closes_in_ns,
-):
+    opens_in_ns: np.ndarray,
+    break_starts_in_ns: np.ndarray,
+    break_ends_in_ns: np.ndarray,
+    closes_in_ns: np.ndarray,
+) -> np.ndarray:
     """
     Given arrays of opens and closes (in nanoseconds) and optionally
     break_starts and break ends, return an array of each minute between the
@@ -71,3 +82,98 @@ def compute_all_minutes(
             )
     out = np.concatenate(pieces).view("datetime64[ns]")
     return out
+
+
+def parse_session(
+    session: Session,
+    param_name: str | None = None,
+    calendar: ExchangeCalendar | None = None,
+    strict: bool = True,
+) -> pd.Timestamp:
+    """Parse input intended to represent a session label.
+
+    Parameters
+    ----------
+    session :
+        Input to be parsed to session label. Must be valid input to
+        pd.Timestamp and have a time component of 00:00.
+
+    param_name : optional
+        Name of a parameter that was to receive a session label. If passed
+        then error message will make reference to the parameter by name.
+
+    calendar : optional
+        Calendar against which to evaluate `session`. Not required
+        if `strict` is False.
+
+    strict : default: True
+        Determines behaviour if `session` parses as UTC midnight although
+        is not a session of `calendar`.
+            True - raise NotSessionError.
+            False - return UTC midnight pd.Timestamp that does not
+            represent a session.
+
+    Returns
+    -------
+    pd.Timestamp
+        pd.Timestamp (UTC with time component of 00:00). If `strict` True
+        then return will represent a session of `calendar`.
+
+    Raises
+    ------
+    TypeError
+        If `session` is not of type pd.Timestamp | str | int | float |
+            datetime.datetime.
+
+    ValueError
+        If `session` is not an acceptable single-argument input to
+        pd.Timestamp.
+
+        If `session` time component is not 00:00.
+
+        If `session` is timezone aware and timezone is not UTC.
+
+    exchange_calendars.errors.NotSessionError
+        If `strict` True and `session` parses to a valid representation of
+        a session label although it is not a session of `calendar`.
+    """
+    if calendar is None and strict is True:
+        raise ValueError("`calendar` must be passed if `strict` True.")
+
+    try:
+        ts = pd.Timestamp(session)
+    except Exception as e:
+        insert = (
+            "received" if param_name is None else f"'{param_name}' received as"
+        )
+        msg = (
+            "A session label must be passed as a pd.Timestamp or a valid"
+            " single-argument input to pd.Timestamp although"
+            f" {insert} '{session}'."
+        )
+        if isinstance(e, TypeError):
+            raise TypeError(msg) from e
+        else:
+            raise ValueError(msg) from e
+
+    if not (ts.tz is None or ts.tz.zone == "UTC"):
+        insert = " " if param_name is None else f" '{param_name}' "
+        raise ValueError(
+            "A session label must be timezone naive or have timezone"
+            f" as 'UTC', although{insert}parsed as '{ts}'."
+        )
+
+    if not ts == ts.normalize():
+        insert = " " if param_name is None else f" '{param_name}' "
+        raise ValueError(
+            "A session label must have a time component of 00:00"
+            f" although{insert}parsed as '{ts}'."
+        )
+
+    if ts.tz is None:
+        ts = ts.tz_localize("UTC")
+
+    if not strict or calendar.is_session(ts):
+        return ts
+    else:
+        raise errors.NotSessionError(calendar, ts, param_name)

exchange_calendars/errors.py

@@ -12,8 +12,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+import typing
+import pandas as pd
+
 from exchange_calendars.utils.memoize import lazyval
 
+if typing.TYPE_CHECKING:
+    from exchange_calendars import ExchangeCalendar
+
 
 class CalendarError(Exception):
     msg = None
@@ -80,3 +87,55 @@ class ScheduleFunctionInvalidCalendar(CalendarError):
         "Invalid calendar '{given_calendar}' passed to schedule_function. "
         "Allowed options are {allowed_calendars}."
     )
+
+
+class NotSessionError(ValueError):
+    """
+    Raised if parameter expecting a session label receives input that
+    parses correctly (UTC midnight) although is not a session.
+
+    Parameters
+    ----------
+    calendar :
+        Calendar for which `ts` assumed as a session.
+
+    ts :
+        Timestamp assumed as a session.
+
+    param_name : optional
+        Name of a parameter that was to receive a session label. If passed
+        then error message will make reference to the parameter by name.
+    """
+
+    def __init__(
+        self,
+        calendar: ExchangeCalendar,
+        ts: pd.Timestamp,
+        param_name: str | None = None,
+    ):
+        self.calendar = calendar
+        self.ts = ts
+        self.param_name = param_name
+
+    def __str__(self) -> str:
+        if self.param_name is not None:
+            msg = (
+                f"Parameter `{self.param_name}` takes a session label"
+                f" although received input that parsed to '{self.ts}' which"
+            )
+        else:
+            msg = f"'{self.ts}'"
+
+        if self.ts < self.calendar.first_session:
+            msg += (
+                " is earlier than the first session of calendar"
+                f" '{self.calendar.name}' ('{self.calendar.first_session}')."
+            )
+        elif self.ts > self.calendar.last_session:
+            msg += (
+                " is later than the last session of calendar"
+                f" '{self.calendar.name}' ('{self.calendar.last_session}')."
+            )
+        else:
+            msg += f" is not a session of calendar '{self.calendar.name}'."
+        return msg