Merge pull request #41 from MarcoGorelli/is-workday

finish is workday

.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/[email protected]
+      - 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
-          pip install polars-xdt --find-links dist --force-reinstall
-          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
-          pip install polars-xdt --find-links dist --force-reinstall
-          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
-          pip install polars-xdt --find-links dist --force-reinstall
-          pip install pytest
-          pytest
 
   sdist:
     runs-on: ubuntu-latest

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]
 

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

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
+numpydoc_show_class_members = False

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,

requirements.txt

@@ -6,3 +6,5 @@ pandas
 pytest
 holidays
 mypy
+black
+ruff

src/expressions.rs

@@ -18,6 +18,7 @@ fn bday_output(input_fields: &[Field]) -> PolarsResult<Field> {
 }
 
 #[polars_expr(output_type_func=bday_output)]
+
 fn advance_n_days(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult<Series> {
     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<Series> {
     let begin_dates = &inputs[0];
@@ -37,6 +39,7 @@ fn sub(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult<Series> {
     impl_sub(begin_dates, end_dates, &weekmask, holidays)
 }
 
+
 #[polars_expr(output_type=Boolean)]
 fn is_workday(inputs: &[Series], kwargs: BusinessDayKwargs) -> PolarsResult<Series> {
     let dates = &inputs[0];

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<Series> {
-    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())
+    }
 }

tests/test_date_range.py

@@ -88,4 +88,3 @@ def test_eager_custom_holiday() -> None:
         ],
     )
     assert_series_equal(result, expected)
-

tests/test_is_busday.py

@@ -50,3 +50,31 @@ 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