Repository: ericmjl/pyjanitor (python)

Issue:
[ENH] support .query()-style strings in df.update_where()
Given the `update_where` example:

```python
# The dataframe must be assigned to a variable first.
data = {
    "a": [1, 2, 3, 4],
    "b": [5, 6, 7, 8],
    "c": [0, 0, 0, 0]
}
df = pd.DataFrame(data)
df = (
    df
    .update_where(
        condition=(df['a'] > 2) & (df['b'] < 8),
        target_column_name='c',
        target_val=10)
    )
```

It would be nice to be able to use `.query`'s style instead to avoid having to use `df.pipe()` if `df` has been modified somewhere along your method chain block:

```python
df = (
    pd.DataFrame(data)
    .update_where(
        condition='a > 2 and b < 8',
        target_column_name='c',
        target_val=10)
    )
)
```

Pull request:
[ENH] updates to update_where and find_replace functions
<!-- Thank you for your PR!

BEFORE YOU CONTINUE! Please add the appropriate three-letter abbreviation to your title.

The abbreviations can be:
- [DOC]: Documentation fixes.
- [ENH]: Code contributions and new features.
- [TST]: Test-related contributions.
- [INF]: Infrastructure-related contributions.

Also, do not forget to tag the relevant issue here as well.

Finally, as commits come in, don't forget to regularly rebase!
-->

# PR Description

Please describe the changes proposed in the pull request:

- Changed update_where conditions to query style mode
- Removed find_replace dependency on update_where function
-

<!-- Doing so provides maintainers with context on what the PR is, and can help us more effectively review your PR. -->

<!-- Please also identify below which issue that has been raised that you are going to close. -->

**This PR resolves #663 .**

<!-- As you go down the PR template, please feel free to delete sections that are irrelevant. -->

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [x] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.rst`.
<!-- We'd like to acknowledge your contributions! -->
3. [ ] Add a line to `CHANGELOG.rst` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

## Quick Check

To do a very quick check that everything is correct, follow these steps below:

- [x] Run the command `make check` from pyjanitor's top-level directory. This will automatically run:
    - black formatting
    - flake8 checking
    - running the test suite
    - docs build

Once done, please check off the check-box above.

If `make check` does not work for you, you can execute the commands listed in the Makefile individually.

## Code Changes

<!-- If you have not made code changes, please feel free to delete this section. -->

If you are adding code changes, please ensure the following:

- [x] Ensure that you have added tests.
- [x] Run all tests (`$ pytest .`) locally on your machine.
    - [ ] Check to ensure that test coverage covers the lines of code that you have added.
    - [ ] Ensure that all tests pass.

## Documentation Changes

<!-- If you have not made documentation changes, please feel free to delete this section. -->

If you are adding documentation changes, please ensure the following:

- [ ] Build the docs locally.
- [ ] View the docs to check that it renders correctly.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7396878..fe5250a 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,9 @@
 new version (on deck)
 =====================
-
+- [ENH]: Upgraded ``update_where`` function to use either the pandas query style,
+  or boolean indexing via the ``loc`` method.
+  Also updated ``find_replace`` function to use the ``loc`` method directly,
+  instead of routing it through the ``update_where`` function. @samukweku
 
 v0.20.7
 =======
diff --git a/janitor/functions.py b/janitor/functions.py
index bfe0547..7f8f3fa 100644
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -3099,7 +3099,7 @@ def _find_replace(
     if match.lower() == "regex":
         for k, v in mapper.items():
             condition = df[column_name].str.contains(k, regex=True)
-            df = df.update_where(condition, column_name, v)
+            df.loc[condition, column_name] = v
     return df
 
 
@@ -3111,7 +3111,8 @@ def update_where(
     target_column_name: Hashable,
     target_val: Any,
 ) -> pd.DataFrame:
-    """Add multiple conditions to update a column in the dataframe.
+    """
+    Add multiple conditions to update a column in the dataframe.
 
     This method mutates the original DataFrame.
 
@@ -3129,7 +3130,7 @@ def update_where(
         df = (
             df
             .update_where(
-                condition=(df['a'] > 2) & (df['b'] < 8),
+                condition=("a > 2 and b < 8",
                 target_column_name='c',
                 target_val=10)
             )
@@ -3140,17 +3141,26 @@ def update_where(
         # 4 8  0
 
     :param df: The pandas DataFrame object.
-    :param conditions: conditions used to update a target column and target
-        value
+    :param conditions: Conditions used to update a target column
+        and target value.
     :param target_column_name: Column to be updated. If column does not exist
         in dataframe, a new column will be created; note that entries that do
         not get set in the new column will be null.
     :param target_val: Value to be updated
     :returns: An updated pandas DataFrame.
-    :raises: IndexError if **conditions** does not have the same length as
-        **df**.
+    :raises: IndexError if ``conditions`` does not have the same length as
+        ``df``.
+    :raises: TypeError if ``conditions`` is not a pandas-compatible string
+        query.
     """
-    df.loc[conditions, target_column_name] = target_val
+
+    # use query mode if a string expression is passed
+    if isinstance(conditions, str):
+        conditions_index = df.query(conditions).index
+    else:
+        conditions_index = df.loc[conditions].index
+    df.loc[conditions_index, target_column_name] = target_val
+
     return df
 
 


Test patch:
diff --git a/tests/functions/test_update_where.py b/tests/functions/test_update_where.py
index ddcfe1e..5b6263a 100644
--- a/tests/functions/test_update_where.py
+++ b/tests/functions/test_update_where.py
@@ -1,5 +1,8 @@
 import pandas as pd
 import pytest
+from pandas.testing import assert_frame_equal
+
+from janitor.functions import update_where
 
 
 @pytest.mark.functions
@@ -16,3 +19,18 @@ def test_update_where(dataframe):
         ),
         dataframe.replace("Cambridge", "Durham"),
     )
+
+
+def test_update_where_query():
+    """Test that function works with pandas query-style string expression."""
+    df = pd.DataFrame(
+        {"a": [1, 2, 3, 4], "b": [5, 6, 7, 8], "c": [0, 0, 0, 0]}
+    )
+    expected = pd.DataFrame(
+        {"a": [1, 2, 3, 4], "b": [5, 6, 7, 8], "c": [0, 0, 10, 0]}
+    )
+    result = update_where(
+        df, conditions="a > 2 and b < 8", target_column_name="c", target_val=10
+    )
+
+    assert_frame_equal(result, expected)