diff --git a/src/polars_business/polars_business/__init__.py b/src/polars_business/polars_business/__init__.py index 1c0a0af..7d7b923 100644 --- a/src/polars_business/polars_business/__init__.py +++ b/src/polars_business/polars_business/__init__.py @@ -12,11 +12,7 @@ class BusinessDayTools: def __init__(self, expr: pl.Expr): self._expr = expr.cast(pl.Int32) - - def advance_n_days(self, - n, - holidays = None - ) -> pl.Expr: + def advance_n_days(self, n, holidays=None) -> pl.Expr: # if not (isinstance(n, int) and n > 0): # raise ValueError("only positive integers are currently supported for `n`") if holidays is None: @@ -24,18 +20,14 @@ def advance_n_days(self, lib=lib, symbol="advance_n_days", is_elementwise=True, - args = [ + args=[ n, - ], + ], ) else: return self._expr._register_plugin( lib=lib, symbol="advance_n_days", is_elementwise=True, - args = [ - n, - pl.Series([list(set(holidays))]).cast(pl.List(pl.Int32)) - ], + args=[n, pl.Series([list(set(holidays))]).cast(pl.List(pl.Int32))], ) - diff --git a/src/run.py b/src/run.py index 77f4661..ac03c1c 100644 --- a/src/run.py +++ b/src/run.py @@ -3,6 +3,13 @@ from datetime import date, datetime import numpy as np import holidays -hols = holidays.country_holidays('UK', years=[2020]) -print(df.with_columns(pl.col('ts').business.advance_n_days(5, holidays=[date(2000, 1, 1), date(2000, 5, 1)]))) \ No newline at end of file +hols = holidays.country_holidays("UK", years=[2020]) + +print( + df.with_columns( + pl.col("ts").business.advance_n_days( + 5, holidays=[date(2000, 1, 1), date(2000, 5, 1)] + ) + ) +)