diff --git a/README.md b/README.md index a9541f2..965379c 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,11 @@ Then, you'll need to install `polars-business`. Currently, you can do this via P $ pip install polars-business ``` +To use it, you'll need to `import polars_business`, and then you'll be a `.business` accessor +on your expressions! + +Currently there's only a single function: `advance_n_days`. + Example ------- @@ -35,3 +40,22 @@ df = pl.DataFrame({ print(df.with_columns(dates_shifted=pl.col('dates').business.advance_n_days(n=5))) ``` + +Note +---- +Currently, only `pl.Date` datatype is supported. + +What to expected +---------------- +The following will hopefully come relatively soon: +- support for `Datetime`s +- support for custom holiday calendars +- support for rolling forwards/backwards to the next + valid business date (if not already on one) + +Ideas for future development: +- business date range +- support for custom mask + + +Currently there's only a single function: `advance_n_days`. \ No newline at end of file diff --git a/src/polars_business/polars_business/__init__.py b/src/polars_business/polars_business/__init__.py index d847cd6..1638120 100644 --- a/src/polars_business/polars_business/__init__.py +++ b/src/polars_business/polars_business/__init__.py @@ -4,7 +4,7 @@ lib = _get_shared_lib_location(__file__) -__version__ = "0.1.6" +__version__ = "0.1.7" @pl.api.register_expr_namespace("business") @@ -14,8 +14,8 @@ def __init__(self, expr: pl.Expr): def advance_n_days(self, n) -> pl.Expr: - if not (isinstance(n, int) and n > 0): - raise ValueError("only positive integers are currently supported for `n`") + # if not (isinstance(n, int) and n > 0): + # raise ValueError("only positive integers are currently supported for `n`") return self._expr._register_plugin( lib=lib, diff --git a/src/polars_business/pyproject.toml b/src/polars_business/pyproject.toml index 87263d5..0b1bdc0 100644 --- a/src/polars_business/pyproject.toml +++ b/src/polars_business/pyproject.toml @@ -10,7 +10,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -version = "0.1.6" +version = "0.1.7" authors = [ { name="Marco Gorelli", email="33491632+MarcoGorelli@users.noreply.github.com" }, ] diff --git a/src/polars_business/src/expressions.rs b/src/polars_business/src/expressions.rs index e85b6b1..79adb27 100644 --- a/src/polars_business/src/expressions.rs +++ b/src/polars_business/src/expressions.rs @@ -6,23 +6,23 @@ fn advance_n_days(inputs: &[Series]) -> PolarsResult { let ca = inputs[0].i32()?; let n = inputs[1].i32()?.get(0).unwrap(); - let out = ca.apply_values( - |mut x|{ - let mut weekday = (x - 4) % 7; + let out = ca.try_apply( + |x|{ + let weekday = (x - 4) % 7; - // If on weekend, roll backwards to previous - // valid date (following pandas here). if weekday == 5 { - x -= 1; - weekday = 4; + polars_bail!(ComputeError: "Saturday is not a business date, cannot advance. `roll` argument coming soon.") } else if weekday == 6 { - x -= 2; - weekday = 4; + polars_bail!(ComputeError: "Sunday is not a business date, cannot advance. `roll` argument coming soon.") } - let n_days = n + (n + weekday) / 5 * 2; - x + n_days + let n_days = if n >= 0 { + n + (n + weekday) / 5 * 2 + } else { + -(-n + (-n + 4-weekday) / 5 * 2) + }; + Ok(x + n_days) } - ); + )?; Ok(out.cast(&DataType::Date).unwrap().into_series()) } diff --git a/src/run.py b/src/run.py index 20528ab..e5e90da 100644 --- a/src/run.py +++ b/src/run.py @@ -8,8 +8,8 @@ }) df = df.filter(pl.col('dates').dt.weekday() <6) -print(df.head().with_columns(dates_shifted=pl.col('dates').business.advance_n_days(n=3))[:5]) -print(df.head().with_columns(dates_shifted=pl.Series(np.busday_offset(df.head()['dates'], 3)))[:5]) +print(df.head().with_columns(dates_shifted=pl.col('dates').business.advance_n_days(n=-3))[:5]) +print(df.head().with_columns(dates_shifted=pl.Series(np.busday_offset(df.head()['dates'], -3)))[:5]) import pandas as pd dfpd = df.to_pandas() diff --git a/tests/test_business_offsets.py b/tests/test_business_offsets.py index fd60b5b..d73c3e8 100644 --- a/tests/test_business_offsets.py +++ b/tests/test_business_offsets.py @@ -6,12 +6,12 @@ from hypothesis import given, assume import polars as pl -from polars_business import BusinessDayTools +import polars_business @given( date=st.dates(min_value=dt.date(2000, 1, 1), max_value=dt.date(9999, 12, 31)), - n=st.integers(min_value=1, max_value=30), + n=st.integers(min_value=-30, max_value=30), ) def test_against_np_busday_offset(date: dt.date, n: int) -> None: assume(date.weekday() < 5) @@ -22,9 +22,11 @@ def test_against_np_busday_offset(date: dt.date, n: int) -> None: @given( date=st.dates(min_value=dt.date(2000, 1, 1), max_value=dt.date(9999, 12, 31)), - n=st.integers(min_value=1, max_value=30), + n=st.integers(min_value=-30, max_value=30), ) def test_against_pandas_bday_offset(date: dt.date, n: int) -> None: + assume(date.weekday() < 5) + result = pl.DataFrame({'ts': [date]}).select(pl.col('ts').business.advance_n_days(n=n))['ts'].item() result = pl.DataFrame({'ts': [date]}).select(pl.col('ts').business.advance_n_days(n=n))['ts'].item() expected = pd.Timestamp(date) + pd.tseries.offsets.BusinessDay(n) assert pd.Timestamp(result) == expected