64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
"""Tools for working with the Polars duration string language."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import TYPE_CHECKING
|
|
from typing import Literal
|
|
from typing import cast
|
|
from typing import get_args
|
|
|
|
if TYPE_CHECKING:
|
|
from typing_extensions import TypeAlias
|
|
|
|
__all__ = ["IntervalUnit", "parse_interval_string"]
|
|
|
|
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.
|
|
"""
|
|
|
|
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])
|
|
|
|
|
|
def parse_interval_string(every: str) -> tuple[int, IntervalUnit]:
|
|
"""Parse a string like "1d", "2h", "3m" into a tuple of (number, unit).
|
|
|
|
Returns:
|
|
A tuple of multiple and unit parsed from the interval string.
|
|
"""
|
|
if match := PATTERN_INTERVAL.match(every):
|
|
multiple = int(match["multiple"])
|
|
unit = cast("IntervalUnit", match["unit"])
|
|
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 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)
|