From 4459b9121bbf4e9d5b9f07dc9d5bc90a327003f4 Mon Sep 17 00:00:00 2001 From: MarcoGorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Fri, 5 Jan 2024 20:30:49 +0000 Subject: [PATCH 1/5] finish is workday --- docs/API.rst | 1 + src/expressions.rs | 3 +++ src/is_workday.rs | 43 ++++++++++++++++++++++++++++++----------- tests/test_is_busday.py | 25 ++++++++++++++++++++++++ 4 files changed, 61 insertions(+), 11 deletions(-) diff --git a/docs/API.rst b/docs/API.rst index 8cd9954..64e548e 100644 --- a/docs/API.rst +++ b/docs/API.rst @@ -7,3 +7,4 @@ API polars_xdt.date_range polars_xdt.workday_count polars_xdt.ExprXDTNamespace.offset_by + polars_xdt.ExprXDTNamespace.is_workday diff --git a/src/expressions.rs b/src/expressions.rs index 15dabe7..2874e5b 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -18,6 +18,7 @@ fn bday_output(input_fields: &[Field]) -> PolarsResult { } #[polars_expr(output_type_func=bday_output)] + fn advance_n_days(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult { let s = &inputs[0]; let n = &inputs[1].cast(&DataType::Int32)?; @@ -28,6 +29,7 @@ fn advance_n_days(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult< impl_advance_n_days(s, n, holidays, &weekmask, &roll) } + #[polars_expr(output_type=Int32)] fn sub(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult { let begin_dates = &inputs[0]; @@ -37,6 +39,7 @@ fn sub(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult { impl_sub(begin_dates, end_dates, &weekmask, holidays) } + #[polars_expr(output_type=Boolean)] fn is_workday(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult { let dates = &inputs[0]; diff --git a/src/is_workday.rs b/src/is_workday.rs index 22ac8fc..c8f215a 100644 --- a/src/is_workday.rs +++ b/src/is_workday.rs @@ -2,30 +2,51 @@ use crate::business_days::weekday; use polars::prelude::*; use pyo3_polars::export::polars_core::utils::{arrow::array::BooleanArray, CustomIterTools}; +fn is_workday_date(date: i32, weekmask: &[bool; 7], holidays: &[i32]) -> bool { + return unsafe { *weekmask.get_unchecked(weekday(date) as usize - 1) } + && (!holidays.contains(&date)); +} + pub(crate) fn impl_is_workday( dates: &Series, weekmask: &[bool; 7], holidays: &[i32], ) -> PolarsResult { - let out = match dates.dtype() { + match dates.dtype() { DataType::Date => { - let dates = dates.date()?; - dates.downcast_iter().map(|arr| -> BooleanArray { + let ca = dates.date()?; + let chunks = ca.downcast_iter().map(|arr| -> BooleanArray { + arr.into_iter() + .map(|date| date.map(|date| is_workday_date(*date, weekmask, holidays))) + .collect_trusted() + }); + Ok(BooleanChunked::from_chunk_iter(ca.name(), chunks).into_series()) + } + DataType::Datetime(time_unit, _time_zone) => { + let multiplier = match time_unit { + TimeUnit::Milliseconds => 60 * 60 * 24 * 1_000, + TimeUnit::Microseconds => 60 * 60 * 24 * 1_000_000, + TimeUnit::Nanoseconds => 60 * 60 * 24 * 1_000_000_000, + }; + let ca = &polars_ops::prelude::replace_time_zone( + dates.datetime()?, + None, + &StringChunked::from_iter(std::iter::once("raise")), + )?; + let chunks = ca.downcast_iter().map(|arr| -> BooleanArray { arr.into_iter() .map(|date| { - date.map( - |date| - unsafe { *weekmask.get_unchecked(weekday(*date) as usize - 1) } - && (!holidays.contains(date)) - ) + date.map(|date| { + is_workday_date((*date / multiplier) as i32, weekmask, holidays) + }) }) .collect_trusted() - }) + }); + Ok(BooleanChunked::from_chunk_iter(ca.name(), chunks).into_series()) } _ => { polars_bail!(InvalidOperation: "polars_xdt is_workday currently only works on Date type. \ For now, please cast to Date first.") } - }; - Ok(BooleanChunked::from_chunk_iter(dates.name(), out).into_series()) + } } diff --git a/tests/test_is_busday.py b/tests/test_is_busday.py index d142abd..2f348a8 100644 --- a/tests/test_is_busday.py +++ b/tests/test_is_busday.py @@ -50,3 +50,28 @@ def test_against_np_is_busday( weekmask = [0 if reverse_mapping[i] in weekend else 1 for i in range(1, 8)] expected = np.is_busday(date, weekmask=weekmask, holidays=holidays) assert result == expected + +@given( + datetime=st.datetimes(min_value=dt.datetime(2000, 1, 1), max_value=dt.datetime(2000, 12, 31)), + weekend=st.lists( + st.sampled_from(["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]), + min_size=0, + max_size=6, + unique=True, + ), + holidays=st.lists( + st.dates(min_value=dt.date(2000, 1, 1), max_value=dt.date(2000, 12, 31)), + min_size=1, + max_size=300, + ), +) +def test_against_np_is_busday_datetime( + datetime: dt.datetime, + weekend: list[str], + holidays: list[dt.date], +) -> None: + result = get_result(datetime, weekend=weekend, holidays=holidays) + weekmask = [0 if reverse_mapping[i] in weekend else 1 for i in range(1, 8)] + date = dt.date(datetime.year, datetime.month, datetime.day) + expected = np.is_busday(date, weekmask=weekmask, holidays=holidays) + assert result == expected From a562201612151eb4aab7e00dc2d341fe1afb8ac2 Mon Sep 17 00:00:00 2001 From: MarcoGorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Fri, 5 Jan 2024 20:39:24 +0000 Subject: [PATCH 2/5] fix ci? --- .github/workflows/CI.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f46bec7..474780a 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -48,7 +48,6 @@ jobs: run: | set -e pip install -r requirements.txt - pip install polars-xdt --find-links dist --force-reinstall pip install pytest pytest @@ -80,7 +79,6 @@ jobs: run: | set -e pip install -r requirements.txt - pip install polars-xdt --find-links dist --force-reinstall pip install pytest pytest @@ -111,7 +109,6 @@ jobs: run: | set -e pip install -r requirements.txt - pip install polars-xdt --find-links dist --force-reinstall pip install pytest pytest From aa924990b1c574c3ee15ea80d071106a322cc532 Mon Sep 17 00:00:00 2001 From: MarcoGorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Fri, 5 Jan 2024 20:44:45 +0000 Subject: [PATCH 3/5] fix ci? --- .github/workflows/CI.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 474780a..9b6cc8a 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -48,6 +48,7 @@ jobs: run: | set -e pip install -r requirements.txt + make install pip install pytest pytest @@ -79,6 +80,7 @@ jobs: run: | set -e pip install -r requirements.txt + make install pip install pytest pytest @@ -109,6 +111,7 @@ jobs: run: | set -e pip install -r requirements.txt + make install pip install pytest pytest From 58fed0e4b068f19fbdcdd55768eb91625738dcef Mon Sep 17 00:00:00 2001 From: MarcoGorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Sat, 6 Jan 2024 08:59:58 +0000 Subject: [PATCH 4/5] separate tests and wheels jobs --- .github/workflows/CI.yml | 52 ++++++++++++++++++---------------------- 1 file changed, 23 insertions(+), 29 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 9b6cc8a..e8fdd08 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -19,17 +19,38 @@ permissions: contents: read jobs: - linux: + + linux_tests: runs-on: ubuntu-latest strategy: matrix: + target: [x86_64] python-version: ["3.8", "3.9", "3.10", "3.11"] - target: [x86_64, x86, aarch64, armv7] steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} + + - name: Set up Rust + run: rustup show + - uses: mozilla-actions/sccache-action@v0.0.3 + - run: make install + - run: venv/bin/python -m pytest tests + - run: venv/bin/python -m black --check . + - run: venv/bin/python -m mypy polars_xdt tests + - run: venv/bin/python -m ruff . + + linux: + runs-on: ubuntu-latest + strategy: + matrix: + target: [x86_64, x86, aarch64, armv7] + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: '3.10' - name: Build wheels uses: PyO3/maturin-action@v1 with: @@ -42,15 +63,6 @@ jobs: with: name: wheels path: dist - - name: pytest - if: ${{ startsWith(matrix.target, 'x86_64') }} - shell: bash - run: | - set -e - pip install -r requirements.txt - make install - pip install pytest - pytest windows: runs-on: windows-latest @@ -74,15 +86,6 @@ jobs: with: name: wheels path: dist - - name: pytest - if: ${{ !startsWith(matrix.target, 'aarch64') && matrix.target != 'x86' }} - shell: bash - run: | - set -e - pip install -r requirements.txt - make install - pip install pytest - pytest macos: runs-on: macos-latest @@ -105,15 +108,6 @@ jobs: with: name: wheels path: dist - - name: pytest - if: ${{ !startsWith(matrix.target, 'aarch64') }} - shell: bash - run: | - set -e - pip install -r requirements.txt - make install - pip install pytest - pytest sdist: runs-on: ubuntu-latest From 7ab073f2db3a96af44e3bda620ee5f2144d2f849 Mon Sep 17 00:00:00 2001 From: MarcoGorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Sat, 6 Jan 2024 09:02:10 +0000 Subject: [PATCH 5/5] lint --- .github/workflows/CI.yml | 2 +- bump_version.py | 16 +++++++++------- docs/conf.py | 31 +++++++++++++++---------------- polars_xdt/__init__.py | 12 +++--------- requirements.txt | 2 ++ tests/test_date_range.py | 1 - tests/test_is_busday.py | 5 ++++- 7 files changed, 34 insertions(+), 35 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index e8fdd08..188fdbd 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -38,7 +38,7 @@ jobs: - run: make install - run: venv/bin/python -m pytest tests - run: venv/bin/python -m black --check . - - run: venv/bin/python -m mypy polars_xdt tests + - run: venv/bin/python -m mypy polars_xdt #tests - run: venv/bin/python -m ruff . linux: diff --git a/bump_version.py b/bump_version.py index 8d437d0..32bbbed 100644 --- a/bump_version.py +++ b/bump_version.py @@ -6,13 +6,15 @@ # Docs are in a a separate repo cause I couldn't figure out # how to deploy from readthedocs -subprocess.run(['make', 'install']) -subprocess.run(['make', 'clean'], cwd='docs') -subprocess.run(['make', 'html'], cwd='docs') -os.system('cp docs/_build/html/* ../docs-polars-xdt/ -r') -subprocess.run(['git', 'add', '.'], cwd='../docs-polars-xdt') -subprocess.run(['git', 'commit', '-m', '\"new version\"', '--allow-empty'], cwd='../docs-polars-xdt') -subprocess.run(['git', 'push', 'origin', 'HEAD'], cwd='../docs-polars-xdt') +subprocess.run(["make", "install"]) +subprocess.run(["make", "clean"], cwd="docs") +subprocess.run(["make", "html"], cwd="docs") +os.system("cp docs/_build/html/* ../docs-polars-xdt/ -r") +subprocess.run(["git", "add", "."], cwd="../docs-polars-xdt") +subprocess.run( + ["git", "commit", "-m", '"new version"', "--allow-empty"], cwd="../docs-polars-xdt" +) +subprocess.run(["git", "push", "origin", "HEAD"], cwd="../docs-polars-xdt") how = sys.argv[1] diff --git a/docs/conf.py b/docs/conf.py index 8916eb7..e74b77e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -5,38 +5,37 @@ # -- Project information ----------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information +import sphinx_autosummary_accessors # noqa: F401 -project = 'polars-xdt' -copyright = '2023, Marco Gorelli' -author = 'Marco Gorelli' +project = "polars-xdt" +copyright = "2023, Marco Gorelli" +author = "Marco Gorelli" -import sphinx_autosummary_accessors # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration extensions = [ - 'sphinx.ext.autodoc', - 'sphinx.ext.autosummary', + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", "sphinx.ext.githubpages", "sphinx.ext.intersphinx", - 'numpydoc', + "numpydoc", "sphinx_copybutton", "sphinx_design", "sphinx_favicon", - ] - -templates_path = ['_templates'] -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] -numpydoc_show_class_members = False +] +templates_path = ["_templates"] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] +numpydoc_show_class_members = False # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output -html_theme = 'alabaster' -html_static_path = ['_static'] +html_theme = "alabaster" +html_static_path = ["_static"] intersphinx_mapping = { "polars": ("https://pola-rs.github.io/polars/py-polars/html/", None), @@ -48,7 +47,7 @@ html_theme = "pydata_sphinx_theme" html_theme_options = { - 'navigation_with_keys': False, + "navigation_with_keys": False, } -numpydoc_show_class_members = False \ No newline at end of file +numpydoc_show_class_members = False diff --git a/polars_xdt/__init__.py b/polars_xdt/__init__.py index 7b0c7b2..72a0d0a 100644 --- a/polars_xdt/__init__.py +++ b/polars_xdt/__init__.py @@ -9,7 +9,7 @@ from polars_xdt.ranges import date_range from polars.type_aliases import PolarsDataType -from typing import Iterable, Literal, Protocol, Sequence, cast, get_args +from typing import Iterable, Literal, Protocol, Sequence, cast if sys.version_info >= (3, 10): from typing import TypeAlias @@ -198,10 +198,7 @@ def sub( holidays_int = [] else: holidays_int = sorted( - { - (holiday - date(1970, 1, 1)).days - for holiday in holidays - } + {(holiday - date(1970, 1, 1)).days for holiday in holidays} ) if isinstance(end_dates, str): end_dates = pl.col(end_dates) @@ -228,10 +225,7 @@ def is_workday( holidays_int = [] else: holidays_int = sorted( - { - (holiday - date(1970, 1, 1)).days - for holiday in holidays - } + {(holiday - date(1970, 1, 1)).days for holiday in holidays} ) result = self._expr.register_plugin( lib=lib, diff --git a/requirements.txt b/requirements.txt index 7c6a19d..3802661 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,5 @@ pandas pytest holidays mypy +black +ruff diff --git a/tests/test_date_range.py b/tests/test_date_range.py index 1cd0962..509def0 100644 --- a/tests/test_date_range.py +++ b/tests/test_date_range.py @@ -88,4 +88,3 @@ def test_eager_custom_holiday() -> None: ], ) assert_series_equal(result, expected) - diff --git a/tests/test_is_busday.py b/tests/test_is_busday.py index 2f348a8..caa8d64 100644 --- a/tests/test_is_busday.py +++ b/tests/test_is_busday.py @@ -51,8 +51,11 @@ def test_against_np_is_busday( expected = np.is_busday(date, weekmask=weekmask, holidays=holidays) assert result == expected + @given( - datetime=st.datetimes(min_value=dt.datetime(2000, 1, 1), max_value=dt.datetime(2000, 12, 31)), + datetime=st.datetimes( + min_value=dt.datetime(2000, 1, 1), max_value=dt.datetime(2000, 12, 31) + ), weekend=st.lists( st.sampled_from(["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]), min_size=0,