diff --git a/README.md b/README.md index c360de9..1955dc1 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,33 @@ shift_excel_dates( ) ``` +### Excluding fixed study dates with `shift_exceptions` + +Some columns contain a mix of patient-specific dates (which should be shifted) and fixed study-wide dates (e.g. an end-of-study date) that must remain unchanged. Use `shift_exceptions` in `sheet_configs` to list any date values that should never be shifted: + +```python +sheet_configs = { + "patients": { + "patient_id_col": "patient_id", + "date_columns": ["last_alive"], + "shift_exceptions": { + "last_alive": ["2024-12-31"], # end-of-study date — never shift + }, + }, +} + +shift_excel_dates( + input_file="input.xlsx", + output_file="output.xlsx", + patient_sheet="patients", + patient_id_col="patient_id", + sheet_configs=sheet_configs, + seed=42, +) +``` + +The exception date strings are parsed with the same flexible parser used for all date values (supports multiple formats and placeholder strings). Exceptions are matched against the parsed date, so `"2024-12-31"` and `"31-12-2024"` both match the same calendar date. + ### Reproducible Shifts with Linking Table To use the same shifts across multiple runs, save and reuse a linking table: @@ -107,6 +134,7 @@ The function accepts the same parameters as `shift_excel_dates` except `date_for - `date_columns`: List of date column names to shift - `header_row`: (Optional) Zero-based row index for the row that contains column names - `skip_rows_after_header`: (Optional) List of zero-based row indices to exclude from data (e.g. a data-type row immediately below the header) + - `shift_exceptions`: (Optional) Dict mapping column names to lists of date strings that should never be shifted (e.g. a fixed end-of-study date). Dates are parsed using the same flexible parser as regular date values. - `patient_header_row`: (Optional) Zero-based header row for the patient sheet (default: 0). If the patient sheet is in `sheet_configs`, that sheet’s `header_row` is used instead. - `patient_skip_rows`: (Optional) Zero-based row indices to exclude from patient data (e.g. a data-type row). If the patient sheet is in `sheet_configs`, that sheet’s `skip_rows_after_header` is used instead. - `min_shift_days` / `max_shift_days`: Range of days to shift (default: -15 to 15) diff --git a/nuh_helper/date_shift/__init__.py b/nuh_helper/date_shift/__init__.py index 41694d2..b4f1796 100644 --- a/nuh_helper/date_shift/__init__.py +++ b/nuh_helper/date_shift/__init__.py @@ -68,16 +68,12 @@ def _get_patient_ids_and_shift_mappings( .unique() .tolist() ) - logger.info( - "Found %d patient(s) in sheet '%s'", len(patient_ids), patient_sheet - ) + logger.info("Found %d patient(s) in sheet '%s'", len(patient_ids), patient_sheet) if linking_table_path and Path(linking_table_path).exists(): logger.info("Loading shift mappings from '%s'", linking_table_path) shift_mappings = mappings.load_shift_mappings(linking_table_path) - shift_mappings = shift_mappings[ - shift_mappings["patient_id"].isin(patient_ids) - ] + shift_mappings = shift_mappings[shift_mappings["patient_id"].isin(patient_ids)] existing_ids = set(shift_mappings["patient_id"]) missing_ids = [pid for pid in patient_ids if pid not in existing_ids] if missing_ids: @@ -88,9 +84,7 @@ def _get_patient_ids_and_shift_mappings( new_shifts = mappings.generate_shift_mappings( missing_ids, min_shift_days, max_shift_days, seed ) - shift_mappings = pd.concat( - [shift_mappings, new_shifts], ignore_index=True - ) + shift_mappings = pd.concat([shift_mappings, new_shifts], ignore_index=True) else: logger.info("Generating shift mappings for %d patient(s)", len(patient_ids)) shift_mappings = mappings.generate_shift_mappings( @@ -112,6 +106,7 @@ def apply_date_shifts( date_columns: list[str], shift_mappings: pd.DataFrame, date_format: str | None = None, + shift_exceptions: dict[str, list[str]] | None = None, ) -> pd.DataFrame: """ Apply date shifts to specified columns in a DataFrame. @@ -124,6 +119,8 @@ def apply_date_shifts( date_format: Optional date format string (e.g., 'YYYY-MM-DD'). Note: This parameter is kept for API compatibility but formatting is applied at the Excel cell level, not in the DataFrame. + shift_exceptions: Optional dict mapping column names to lists of date strings + that should never be shifted (e.g. a fixed end-of-study date). Returns: DataFrame with shifted dates. @@ -141,6 +138,18 @@ def apply_date_shifts( ) ) + # Pre-parse exception dates once per column + parsed_exceptions: dict[str, set[date]] = {} + if shift_exceptions: + for col, exc_values in shift_exceptions.items(): + parsed_set: set[date] = set() + for v in exc_values: + ts = _parse._parse_date_value(v) + if ts is not None: + parsed_set.add(ts.date()) + if parsed_set: + parsed_exceptions[col] = parsed_set + for date_col in date_columns: if date_col not in df.columns: logger.warning( @@ -151,9 +160,7 @@ def apply_date_shifts( # Parse flexible date strings (handles YYYY-DD-MM and placeholders "Unknown") non_null_before = df[date_col].notna().sum() df[date_col] = df[date_col].apply(_parse._parse_date_value) - parse_failures = non_null_before - sum( - x is not None for x in df[date_col] - ) + parse_failures = non_null_before - sum(x is not None for x in df[date_col]) if parse_failures > 0: logger.debug( "Column '%s': %d value(s) could not be parsed as dates", @@ -162,11 +169,14 @@ def apply_date_shifts( ) # Apply shifts + exc_dates = parsed_exceptions.get(date_col, set()) df[date_col] = df.apply( lambda row: ( row[date_col] # noqa: B023 + pd.Timedelta(days=shift_dict.get(row[patient_id_col], 0)) - if row[date_col] is not None and row[patient_id_col] in shift_dict # noqa: B023 + if row[date_col] is not None # noqa: B023 + and row[patient_id_col] in shift_dict # noqa: B023 + and row[date_col].date() not in exc_dates # noqa: B023 else row[date_col] # noqa: B023 ), axis=1, @@ -175,9 +185,7 @@ def apply_date_shifts( # Convert back to date-only format (removes time component) df[date_col] = df[date_col].apply( lambda x: ( - x.date() - if isinstance(x, (pd.Timestamp, datetime, date)) - else None + x.date() if isinstance(x, (pd.Timestamp, datetime, date)) else None ), ) @@ -256,6 +264,7 @@ def shift_excel_dates( header_row = default_header_row sheet_date_columns: list[str] | None = None skip_rows_after_header: list[int] | None = None + sheet_shift_exceptions: dict[str, list[str]] | None = None if sheet_name in sheet_configs: config = sheet_configs[cast(str, sheet_name)] @@ -263,6 +272,7 @@ def shift_excel_dates( date_columns: list[str] = cast(list[str], config["date_columns"]) header_row = cast(int, config.get("header_row", header_row)) skip_rows_after_header = config.get("skip_rows_after_header") + sheet_shift_exceptions = config.get("shift_exceptions") sheet_date_columns = date_columns logger.info( "Shifting %d date column(s) in sheet '%s'", @@ -292,6 +302,7 @@ def shift_excel_dates( date_columns, shift_mappings, date_format=None, + shift_exceptions=sheet_shift_exceptions, ) _excel._write_sheet_with_structure( @@ -397,9 +408,7 @@ def shift_excel_dates_inplace( sheet_patient_id_col: str = cast(str, config["patient_id_col"]) date_columns: list[str] = cast(list[str], config["date_columns"]) header_row: int = cast(int, config.get("header_row", 0)) - skip_rows_after_header: list[int] | None = config.get( - "skip_rows_after_header" - ) + skip_rows_after_header: list[int] | None = config.get("skip_rows_after_header") max_col = ws.max_column or 0 if not max_col: @@ -434,6 +443,21 @@ def shift_excel_dates_inplace( if not date_col_indices: continue + # Pre-parse exception dates once per column + parsed_exceptions: dict[str, set[date]] = {} + shift_exceptions_config: dict[str, list[str]] | None = config.get( + "shift_exceptions" + ) + if shift_exceptions_config: + for exc_col, exc_values in shift_exceptions_config.items(): + parsed_set: set[date] = set() + for v in exc_values: + ts = _parse._parse_date_value(v) + if ts is not None: + parsed_set.add(ts.date()) + if parsed_set: + parsed_exceptions[exc_col] = parsed_set + skip_row_set: set[int] = ( {idx + 1 for idx in skip_rows_after_header} if skip_rows_after_header @@ -455,7 +479,7 @@ def shift_excel_dates_inplace( pid = _parse._normalize_patient_id(pid_cell.value) shift_days = shift_dict.get(pid) if pid is not None else None - for date_col_idx in date_col_indices.values(): + for col_name, date_col_idx in date_col_indices.items(): cell = ws.cell(row=row_idx, column=date_col_idx) original_value = cell.value parsed = _parse._parse_date_value(original_value) @@ -470,6 +494,10 @@ def shift_excel_dates_inplace( if shift_days is None: continue + exc_dates = parsed_exceptions.get(col_name, set()) + if exc_dates and parsed.date() in exc_dates: + continue + shifted = parsed + pd.Timedelta(days=shift_days) if isinstance(original_value, date) and not isinstance( original_value, datetime @@ -498,4 +526,4 @@ def shift_excel_dates_inplace( "apply_date_shifts", "generate_shift_mappings", "load_shift_mappings", -] \ No newline at end of file +] diff --git a/nuh_helper/date_shift/_excel.py b/nuh_helper/date_shift/_excel.py index 43b831c..204b4a5 100644 --- a/nuh_helper/date_shift/_excel.py +++ b/nuh_helper/date_shift/_excel.py @@ -89,40 +89,28 @@ def _read_sheet_with_structure( # Pandas-friendly column names (no None) columns = [] for i, v in enumerate(header_values): - if v is None or ( - isinstance(v, float) and pd.isna(cast("float", v)) - ): + if v is None or (isinstance(v, float) and pd.isna(cast("float", v))): columns.append(f"Unnamed: {i}") else: columns.append(str(v).strip() or f"Unnamed: {i}") if header_row > 0: - description_merged_ranges = _description_merged_ranges( - ws, header_row - ) + description_merged_ranges = _description_merged_ranges(ws, header_row) # Description rows (rows before header) description_rows = ( - full_df.iloc[:header_row].values.tolist() - if header_row > 0 - else [] + full_df.iloc[:header_row].values.tolist() if header_row > 0 else [] ) # Data rows: everything after header, optionally excluding some rows data_block = full_df.iloc[header_row + 1 :] if skip_rows_after_header: # Drop by 0-based index (data_block index = original row) - to_drop = [ - i - for i in data_block.index - if i in skip_rows_after_header - ] + to_drop = [i for i in data_block.index if i in skip_rows_after_header] data_block = data_block.drop(index=to_drop, errors="ignore") data_df = pd.DataFrame(data_block.values, columns=columns) description_df = ( - full_df.iloc[:header_row].copy() - if header_row > 0 - else pd.DataFrame() + full_df.iloc[:header_row].copy() if header_row > 0 else pd.DataFrame() ) wb.close() return ( @@ -136,17 +124,13 @@ def _read_sheet_with_structure( if header_row == 0: description_rows = [] description_df = pd.DataFrame() - data_df = pd.read_excel( - excel_file, sheet_name=sheet_name, header=0 - ) + data_df = pd.read_excel(excel_file, sheet_name=sheet_name, header=0) if skip_rows_after_header: data_df = data_df.drop(index=skip_rows_after_header, errors="ignore") else: description_rows = full_df.iloc[:header_row].values.tolist() description_df = full_df.iloc[:header_row].copy() - data_df = pd.read_excel( - excel_file, sheet_name=sheet_name, header=header_row - ) + data_df = pd.read_excel(excel_file, sheet_name=sheet_name, header=header_row) if skip_rows_after_header: data_df = data_df.drop(index=skip_rows_after_header, errors="ignore") diff --git a/tests/test_date_shift.py b/tests/test_date_shift.py index b524b06..d624a16 100644 --- a/tests/test_date_shift.py +++ b/tests/test_date_shift.py @@ -1,4 +1,5 @@ """Unit tests for nuh_helper.date_shift""" + from datetime import date, datetime from pathlib import Path @@ -211,13 +212,13 @@ def test_empty_when_num_description_rows_zero(self) -> None: class TestShiftExcelDatesWithComplexLayout: - """Integration tests: header_row, skip_rows_after_header, merged cells, + """Integration tests: header_row, skip_rows_after_header, merged cells, patient sheet from config.""" def _make_excel_with_complex_header( self, path: Path, sheet_name: str = "patients" ) -> None: - """Create an xlsx with: row0 merged title, row1 description, row2 header, row3 + """Create an xlsx with: row0 merged title, row1 description, row2 header, row3 data-type, row4+ data.""" wb = Workbook() ws = wb.active @@ -250,7 +251,7 @@ def _make_excel_with_complex_header( def test_header_row_and_skip_rows_after_header_produce_correct_columns( self, tmp_path: Path ) -> None: - """With header_row=2 and skip_rows_after_header=[3], + """With header_row=2 and skip_rows_after_header=[3], column names come from row 2, row 3 excluded.""" xlsx = tmp_path / "in.xlsx" self._make_excel_with_complex_header(xlsx) @@ -282,7 +283,7 @@ def test_patient_sheet_uses_config_header_when_sheet_in_sheet_configs( self, tmp_path: Path ) -> None: """When patient_sheet is in sheet_configs, its header_row is used - (no patient_header_row arg).""" + (no patient_header_row arg).""" xlsx = tmp_path / "in.xlsx" self._make_excel_with_complex_header(xlsx) out = tmp_path / "out.xlsx" @@ -404,3 +405,131 @@ def test_does_not_mutate_input(self) -> None: df, "patient_id", ["visit_date"], self._make_mappings("P001", 10) ) assert df["visit_date"].iloc[0] == original_visit + + +class TestApplyDateShiftsExceptions: + def _make_mappings(self, patient_id: str, shift_days: int) -> pd.DataFrame: + return pd.DataFrame({"patient_id": [patient_id], "shift_days": [shift_days]}) + + def test_exception_date_is_not_shifted(self) -> None: + df = pd.DataFrame({"patient_id": ["P001"], "visit_date": ["2024-12-31"]}) + result = apply_date_shifts( + df, + "patient_id", + ["visit_date"], + self._make_mappings("P001", 10), + shift_exceptions={"visit_date": ["2024-12-31"]}, + ) + assert result["visit_date"].iloc[0] == date(2024, 12, 31) + + def test_non_exception_date_in_same_column_is_shifted(self) -> None: + df = pd.DataFrame( + { + "patient_id": ["P001", "P002"], + "visit_date": ["2024-12-31", "2023-06-01"], + } + ) + mappings = pd.DataFrame( + {"patient_id": ["P001", "P002"], "shift_days": [10, 10]} + ) + result = apply_date_shifts( + df, + "patient_id", + ["visit_date"], + mappings, + shift_exceptions={"visit_date": ["2024-12-31"]}, + ) + assert result["visit_date"].iloc[0] == date(2024, 12, 31) + assert result["visit_date"].iloc[1] == date(2023, 6, 11) + + def test_exception_only_applies_to_specified_column(self) -> None: + df = pd.DataFrame( + { + "patient_id": ["P001"], + "col_a": ["2024-12-31"], + "col_b": ["2024-12-31"], + } + ) + result = apply_date_shifts( + df, + "patient_id", + ["col_a", "col_b"], + self._make_mappings("P001", 5), + shift_exceptions={"col_a": ["2024-12-31"]}, + ) + assert result["col_a"].iloc[0] == date(2024, 12, 31) + assert result["col_b"].iloc[0] == date(2025, 1, 5) + + def test_multiple_exception_dates(self) -> None: + df = pd.DataFrame( + { + "patient_id": ["P001", "P001", "P001"], + "visit_date": ["2024-12-31", "2024-06-30", "2023-01-15"], + } + ) + result = apply_date_shifts( + df, + "patient_id", + ["visit_date"], + self._make_mappings("P001", 7), + shift_exceptions={"visit_date": ["2024-12-31", "2024-06-30"]}, + ) + assert result["visit_date"].iloc[0] == date(2024, 12, 31) + assert result["visit_date"].iloc[1] == date(2024, 6, 30) + assert result["visit_date"].iloc[2] == date(2023, 1, 22) + + def test_no_exceptions_kwarg_shifts_normally(self) -> None: + df = pd.DataFrame({"patient_id": ["P001"], "visit_date": ["2024-12-31"]}) + result = apply_date_shifts( + df, + "patient_id", + ["visit_date"], + self._make_mappings("P001", 10), + shift_exceptions=None, + ) + assert result["visit_date"].iloc[0] == date(2025, 1, 10) + + +class TestShiftExcelDatesExceptionsIntegration: + """Integration tests for shift_exceptions passed via sheet_configs.""" + + def _make_excel(self, path: Path) -> None: + wb = Workbook() + ws = wb.active + assert ws is not None + ws.title = "patients" + ws["A1"] = "patient_id" + ws["B1"] = "last_alive" + ws["A2"] = "P001" + ws["B2"] = "2023-06-01" + ws["A3"] = "P002" + ws["B3"] = "2024-12-31" + wb.save(path) + + def test_exception_date_unchanged_other_dates_shifted(self, tmp_path: Path) -> None: + xlsx = tmp_path / "in.xlsx" + self._make_excel(xlsx) + out = tmp_path / "out.xlsx" + sheet_configs = { + "patients": { + "patient_id_col": "patient_id", + "date_columns": ["last_alive"], + "shift_exceptions": {"last_alive": ["2024-12-31"]}, + } + } + shift_excel_dates( + input_file=str(xlsx), + output_file=str(out), + patient_sheet="patients", + patient_id_col="patient_id", + sheet_configs=sheet_configs, + linking_table_output=str(tmp_path / "linking.csv"), + seed=42, + ) + df = pd.read_excel(str(out), sheet_name="patients") + p001_date = df.loc[df["patient_id"] == "P001", "last_alive"].iloc[0] + p002_date = df.loc[df["patient_id"] == "P002", "last_alive"].iloc[0] + # P002's date is an exception — must be unchanged + assert pd.Timestamp(p002_date).date() == date(2024, 12, 31) + # P001's date must have been shifted (seed=42 gives a non-zero shift) + assert pd.Timestamp(p001_date).date() != date(2023, 6, 1)