1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
|
"""Tools for working with the Polars duration string language."""
from __future__ import annotations
import re
from typing import TYPE_CHECKING, Literal, cast, 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)
|
