simplify

src/Makefile

@@ -2,24 +2,39 @@
 SHELL=/bin/bash
 
 venv:  ## Set up virtual environment
-	python3 -m venv venv
-	venv/bin/pip install -r requirements.txt
+	. ../.venv/bin/activate
 
 install: venv
 	unset CONDA_PREFIX && \
-	source venv/bin/activate && maturin develop -m polars_business/Cargo.toml
+	source ../.venv/bin/activate && maturin develop -m polars_business/Cargo.toml
 
 install-release: venv
 	unset CONDA_PREFIX && \
-	source venv/bin/activate && maturin develop --release -m polars_business/Cargo.toml
+	source ../.venv/bin/activate && maturin develop --release -m polars_business/Cargo.toml
+
+build: venv
+	unset CONDA_PREFIX && \
+	source ../.venv/bin/activate && maturin develop --release -m polars_business/Cargo.toml
 
 clean:
-	-@rm -r venv
+	-@rm -r ../.venv
 	-@cd polars_business && cargo clean
 
 
 run: install
-	source venv/bin/activate && python run.py
+	source ../.venv/bin/activate && python run.py
 
 run-release: install-release
-	source venv/bin/activate && python run.py
+	source ../.venv/bin/activate && python run.py
+
+test: build
+	pytest ../tests
+
+clippy: venv
+	cargo clippy -p polars_business
+
+fmt: venv
+	cargo fmt --all
+	black .
+
+pre-commit: clippy fmt 

src/polars_business/src/expressions.rs

@@ -6,23 +6,53 @@ fn weekday(x: i32) -> i32 {
     (x - 4) % 7
 }
 
-fn calculate_n_days_without_holidays(n: i32, x_weekday: i32) -> i32 {
+fn calculate_n_days_without_holidays(x: i32, n: i32, x_weekday: i32) -> i32 {
     if n >= 0 {
-        n + (n + x_weekday) / 5 * 2
+        // Let's pretend we were starting on a Monday. How many days would we
+        // need to advance?
+        let mut n_days = n + n / 5 * 2 - x_weekday;
+
+        // Right. But we didn't necessarily start on a Monday, we started on
+        // x_weekday. So now, let's advance by x_weekday days, each time
+        // rolling forwards if we need to. x_weekday <= 6 so this loop won't
+        // happen too many times anyway.
+        let mut i = x_weekday;
+        while i > 0 {
+            n_days += 1;
+            let res_weekday = weekday(x + n_days);
+            n_days = roll(n_days, res_weekday);
+            i -= 1;
+        }
+        n_days
     } else {
-        -(-n + (-n + 4 - x_weekday) / 5 * 2)
+        // Reverse logic to the above.
+        // Let's pretend we were starting on a Friday.
+        // How many days to advance?
+        let mut n_days = n + n / 5 * 2 + (4 - x_weekday);
+
+        // Then, adjust
+        let mut i = 4 - x_weekday;
+        while i > 0 {
+            n_days -= 1;
+            let res_weekday = weekday(x + n_days);
+            n_days = roll(n_days, res_weekday);
+            i -= 1;
+        }
+        n_days
     }
 }
 
 fn reduce_vec(vec: &[Option<i32>], x: i32, n_days: i32) -> Vec<Option<i32>> {
     // Each day we skip may be a holiday, and so require skipping an additional day.
     // n_days*2 is an upper-bound.
     if n_days > 0 {
-        vec.iter().copied()
+        vec.iter()
+            .copied()
             .filter(|t| t.map(|t| t >= x && t <= x + n_days * 2).unwrap_or(false))
             .collect()
     } else {
-        vec.iter().copied()
+        vec.iter()
+            .copied()
             .filter(|t| t.map(|t| t <= x && t >= x + n_days * 2).unwrap_or(false))
             .collect()
     }
@@ -63,7 +93,7 @@ fn calculate_n_days(x: i32, n: i32, vec: &Vec<Option<i32>>) -> PolarsResult<i32>
         polars_bail!(ComputeError: "Sunday is not a business date, cannot advance. `roll` argument coming soon.")
     }
 
-    let mut n_days = calculate_n_days_without_holidays(n, x_weekday);
+    let mut n_days = calculate_n_days_without_holidays(x, n, x_weekday);
 
     if !vec.is_empty() {
         let mut myvec: Vec<Option<i32>> = reduce_vec(vec, x, n_days);

src/run.py

@@ -2,14 +2,35 @@
 from polars_business import *
 from datetime import date, datetime
 import numpy as np
-import holidays
 
-hols = holidays.country_holidays("UK", years=[2020])
+df = pl.DataFrame(
+    {
+        "dates": pl.date_range(date(2000, 12, 25), date(2000, 12, 30), eager=True),
+    }
+)
+df = df.filter(pl.col("dates").dt.weekday() < 6)
+df = df.with_columns(start_wday=pl.col("dates").dt.strftime("%a"))
 
 print(
     df.with_columns(
-        pl.col("ts").business.advance_n_days(
-            5, holidays=[date(2000, 1, 1), date(2000, 5, 1)]
+        dates_shifted=pl.col("dates").business.advance_n_days(
+            n=-3,
+            # holidays=[date(2000, 12, 25)]
+        )
+    ).with_columns(end_wday=pl.col("dates_shifted").dt.strftime("%a"))
+)
+print(
+    df.with_columns(
+        dates_shifted=pl.Series(
+            np.busday_offset(
+                df["dates"],
+                -3,
+                #   holidays=[date(2000, 12, 25)]
+            )
         )
     )
 )
+
+# E           date=datetime.date(2000, 4, 24),
+# E           n=5,
+# E           holidays=[datetime.date(2000, 1, 1), datetime.date(2000, 5, 1)],