Skip to content

Commit

Permalink
cache day names
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Apr 9, 2024
1 parent 6dd3036 commit 81cd34a
Showing 1 changed file with 16 additions and 12 deletions.
28 changes: 16 additions & 12 deletions py-polars/polars/functions/business.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import contextlib
import functools
from typing import TYPE_CHECKING, Iterable

from polars._utils.parse_expr_input import parse_as_expression
Expand All @@ -15,15 +16,18 @@
from polars import Expr
from polars.type_aliases import DayOfWeek, IntoExprColumn

DAY_NAMES = (
"Mon",
"Tue",
"Wed",
"Thu",
"Fri",
"Sat",
"Sun",
)

@functools.lru_cache
def _day_names() -> tuple[str, ...]:
return (
"Mon",
"Tue",
"Wed",
"Thu",
"Fri",
"Sat",
"Sun",
)


def _make_week_mask(
Expand All @@ -36,13 +40,13 @@ def _make_week_mask(
else:
weekend_set = set(weekend)
for day in weekend_set:
if day not in DAY_NAMES:
msg = f"Expected one of {DAY_NAMES}, got: {day}"
if day not in _day_names():
msg = f"Expected one of {_day_names()}, got: {day}"
raise ValueError(msg)
return tuple(
[
False if v in weekend else True # noqa: SIM211
for v in DAY_NAMES
for v in _day_names()
]
)

Expand Down

0 comments on commit 81cd34a

Please sign in to comment.