done

lib/python3.11/site-packages/narwhals/_duration.py

@@ -0,0 +1,94 @@
+"""Tools for working with the Polars duration string language."""
+
+from __future__ import annotations
+
+import datetime as dt
+import re
+from typing import TYPE_CHECKING, Literal, cast, get_args
+
+if TYPE_CHECKING:
+    from collections.abc import Container, Mapping
+
+    from typing_extensions import TypeAlias
+
+__all__ = ["IntervalUnit"]
+
+IntervalUnit: TypeAlias = Literal["ns", "us", "ms", "s", "m", "h", "d", "mo", "q", "y"]
+"""A Polars duration string interval unit.
+
+- 'ns': nanosecond.
+- 'us': microsecond.
+- 'ms': millisecond.
+- 's': second.
+- 'm': minute.
+- 'h': hour.
+- 'd': day.
+- 'mo': month.
+- 'q': quarter.
+- 'y': year.
+"""
+TimedeltaKwd: TypeAlias = Literal[
+    "days", "hours", "minutes", "seconds", "milliseconds", "microseconds"
+]
+
+PATTERN_INTERVAL: re.Pattern[str] = re.compile(
+    r"^(?P<multiple>-?\d+)(?P<unit>ns|us|ms|mo|m|s|h|d|q|y)\Z"
+)
+MONTH_MULTIPLES = frozenset([1, 2, 3, 4, 6, 12])
+QUARTER_MULTIPLES = frozenset([1, 2, 4])
+UNIT_TO_TIMEDELTA: Mapping[IntervalUnit, TimedeltaKwd] = {
+    "d": "days",
+    "h": "hours",
+    "m": "minutes",
+    "s": "seconds",
+    "ms": "milliseconds",
+    "us": "microseconds",
+}
+
+
+class Interval:
+    def __init__(self, multiple: int, unit: IntervalUnit, /) -> None:
+        self.multiple: int = multiple
+        self.unit: IntervalUnit = unit
+
+    def to_timedelta(
+        self, *, unsupported: Container[IntervalUnit] = frozenset(("ns", "mo", "q", "y"))
+    ) -> dt.timedelta:
+        if self.unit in unsupported:  # pragma: no cover
+            msg = f"Creating timedelta with {self.unit} unit is not supported."
+            raise NotImplementedError(msg)
+        kwd = UNIT_TO_TIMEDELTA[self.unit]
+        # error: Keywords must be strings (bad mypy)
+        return dt.timedelta(**{kwd: self.multiple})  # type: ignore[misc]
+
+    @classmethod
+    def parse(cls, every: str) -> Interval:
+        multiple, unit = cls._parse(every)
+        if unit == "mo" and multiple not in MONTH_MULTIPLES:
+            msg = f"Only the following multiples are supported for 'mo' unit: {MONTH_MULTIPLES}.\nGot: {multiple}."
+            raise ValueError(msg)
+        if unit == "q" and multiple not in QUARTER_MULTIPLES:
+            msg = f"Only the following multiples are supported for 'q' unit: {QUARTER_MULTIPLES}.\nGot: {multiple}."
+            raise ValueError(msg)
+        if unit == "y" and multiple != 1:
+            msg = (
+                f"Only multiple 1 is currently supported for 'y' unit.\nGot: {multiple}."
+            )
+            raise ValueError(msg)
+        return cls(multiple, unit)
+
+    @classmethod
+    def parse_no_constraints(cls, every: str) -> Interval:
+        return cls(*cls._parse(every))
+
+    @staticmethod
+    def _parse(every: str) -> tuple[int, IntervalUnit]:
+        if match := PATTERN_INTERVAL.match(every):
+            multiple = int(match["multiple"])
+            unit = cast("IntervalUnit", match["unit"])
+            return multiple, unit
+        msg = (
+            f"Invalid `every` string: {every}. Expected string of kind <number><unit>, "
+            f"where 'unit' is one of: {get_args(IntervalUnit)}."
+        )
+        raise ValueError(msg)