From 1f00a61f468450c751b4e3f24a4c7ef46a645194 Mon Sep 17 00:00:00 2001 From: naseem173 Date: Thu, 23 Jul 2026 23:53:18 +0530 Subject: [PATCH] Add flag_in_range function (resolves #708) --- CHANGELOG.md | 1 + janitor/functions/__init__.py | 8 ++- janitor/functions/flag_in_range.py | 77 +++++++++++++++++++++++++++ tests/functions/test_flag_in_range.py | 77 +++++++++++++++++++++++++++ 4 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 janitor/functions/flag_in_range.py create mode 100644 tests/functions/test_flag_in_range.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 037e915b9..9804ba6ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ - [ENH] Improve `polars.complete` and `polars.expand`, avoiding the potentially expensive schema computation on lazy frames. @samukweku - [ENH] Add `strip_whitespace` to `clean_names` function. - Issue #1385 +- [ENH] Added `flag_in_range` function to flag values outside a given range. - Issue #708 @naseem173 ## [v0.32.3] - 2025-12-11 diff --git a/janitor/functions/__init__.py b/janitor/functions/__init__.py index 3a9fa542d..05b05020d 100644 --- a/janitor/functions/__init__.py +++ b/janitor/functions/__init__.py @@ -33,7 +33,11 @@ from .clean_names import clean_names from .coalesce import coalesce from .collapse_levels import collapse_levels -from .compare_df_cols import compare_df_cols, compare_df_cols_same, describe_class +from .compare_df_cols import ( + compare_df_cols, + compare_df_cols_same, + describe_class, +) from .complete import complete from .concatenate_columns import concatenate_columns from .conditional_join import conditional_join, get_join_indices, join_agg @@ -60,6 +64,7 @@ from .fill import fill_direction, fill_empty from .filter import filter_column_isin, filter_date, filter_on, filter_string from .find_replace import find_replace +from .flag_in_range import flag_in_range from .flag_nulls import flag_nulls from .get_dupes import get_dupes from .get_one_to_one import get_one_to_one @@ -170,6 +175,7 @@ "filter_on", "filter_string", "find_replace", + "flag_in_range", "flag_nulls", "get_dupes", "get_one_to_one", diff --git a/janitor/functions/flag_in_range.py b/janitor/functions/flag_in_range.py new file mode 100644 index 000000000..6f1539897 --- /dev/null +++ b/janitor/functions/flag_in_range.py @@ -0,0 +1,77 @@ +"""Implementation source for `flag_in_range`.""" + +from typing import Hashable, Optional, Union + +import numpy as np +import pandas as pd +import pandas_flavor as pf + +from janitor.utils import check_column + + +@pf.register_dataframe_method +def flag_in_range( + df: pd.DataFrame, + column_name: Hashable, + low: Union[int, float], + high: Union[int, float], + inclusive: bool = True, + flag_column_name: Optional[Hashable] = "in_range_flag", +) -> pd.DataFrame: + """Creates a new column to indicate whether values in a column fall + within (or outside) a given range. + + A flag value of `1` indicates the row's value in `column_name` falls + outside the `[low, high]` range; `0` indicates it falls within range, + consistent with the convention used by `flag_nulls`. + + This method does not mutate the original DataFrame. + + Examples: + >>> import pandas as pd + >>> import janitor + >>> df = pd.DataFrame({"a": [1, 5, 10, 15, 20]}) + >>> df.flag_in_range(column_name="a", low=5, high=15) + a in_range_flag + 0 1 1 + 1 5 0 + 2 10 0 + 3 15 0 + 4 20 1 + + Args: + df: Input pandas DataFrame. + column_name: Name of the column to check. + low: Lower bound of the range. + high: Upper bound of the range. + inclusive: Whether the range bounds (`low` and `high`) are + themselves considered "in range". Defaults to True. + flag_column_name: Name for the output flag column. + + Raises: + ValueError: If `column_name` is not present in the DataFrame. + ValueError: If `flag_column_name` is already present in the + DataFrame. + + Returns: + Input dataframe with the range flag column appended. + + + """ + check_column(df, [column_name]) + check_column(df, [flag_column_name], present=False) + + series = df[column_name] + + if inclusive: + in_range = series.between(low, high) + else: + in_range = (series > low) & (series < high) + + out_of_range = np.logical_not(in_range) + + df = df.copy() + df[flag_column_name] = out_of_range.astype(int) + return df diff --git a/tests/functions/test_flag_in_range.py b/tests/functions/test_flag_in_range.py new file mode 100644 index 000000000..0288c3197 --- /dev/null +++ b/tests/functions/test_flag_in_range.py @@ -0,0 +1,77 @@ +"""Tests for `flag_in_range` function.""" + +import pandas as pd +import pytest +from pandas.testing import assert_frame_equal + +from janitor.functions import flag_in_range + + +@pytest.fixture +def range_df(): + """A small DataFrame for range-flagging tests.""" + return pd.DataFrame({"a": [1, 5, 10, 15, 20]}) + + +@pytest.mark.functions +def test_functional_default(range_df): + """Checks default (inclusive) behaviour as a method call.""" + expected = range_df.copy() + expected["in_range_flag"] = [1, 0, 0, 0, 1] + + df = range_df.flag_in_range(column_name="a", low=5, high=15) + + assert_frame_equal(df, expected, check_dtype=False) + + +@pytest.mark.functions +def test_non_method_functional(range_df): + """Checks behaviour when `flag_in_range` is used as a function.""" + expected = range_df.copy() + expected["in_range_flag"] = [1, 0, 0, 0, 1] + + df = flag_in_range(range_df, column_name="a", low=5, high=15) + + assert_frame_equal(df, expected, check_dtype=False) + + +@pytest.mark.functions +def test_exclusive_bounds(range_df): + """Checks that exclusive bounds flag boundary values as out of range.""" + expected = range_df.copy() + expected["in_range_flag"] = [1, 1, 0, 1, 1] + + df = range_df.flag_in_range( + column_name="a", low=5, high=15, inclusive=False + ) + + assert_frame_equal(df, expected, check_dtype=False) + + +@pytest.mark.functions +def test_rename_output_column(range_df): + """Checks output column is renamed when `flag_column_name` is given.""" + expected = range_df.copy() + expected["flag"] = [1, 0, 0, 0, 1] + + df = range_df.flag_in_range( + column_name="a", low=5, high=15, flag_column_name="flag" + ) + + assert_frame_equal(df, expected, check_dtype=False) + + +@pytest.mark.functions +def test_fail_column_not_in_df(range_df): + """Checks ValueError is raised when `column_name` is not in df.""" + with pytest.raises(ValueError): + range_df.flag_in_range(column_name="z", low=5, high=15) + + +@pytest.mark.functions +def test_fail_flag_column_exists(range_df): + """Checks ValueError is raised when `flag_column_name` already exists.""" + with pytest.raises(ValueError): + range_df.flag_in_range( + column_name="a", low=5, high=15, flag_column_name="a" + )