Repository: pyjanitor-devs/pyjanitor

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)
    )
)
```


Discussion:
Hi @zbarry , if it's ok with you, I would love to work on this.
That’d be great. Thanks!
@zbarry Cool. I'll start off on it. Is it required that you assign it to me?
Good point; I went ahead and assigned it.

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- 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
--- 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)