ENH: add regex parameter to .loc

Author: topper-123

doc/source/user_guide/indexing.rst

@@ -573,6 +573,38 @@ without using a temporary variable.
    (bb.groupby(['year', 'team']).sum()
       .loc[lambda df: df.r > 100])
 
+
+.. _indexing.selecting_with_regex:
+
+Selection by regular expression
+-------------------------------
+
+.. versionadded:: 0.25.0
+
+it is possible to call :attr:`~DataFrame.loc` with parameter ``regex=True`` to select by
+row/columns axis labels that match a regular expression pattern.
+
+.. ipython:: python
+
+    df = pd.DataFrame(1, index=["A", "AB", "BC", "CD"], columns=["CD", "BC", "AB", "A"])
+    df
+
+    df.loc(regex=True)["B", "B"]
+    df.loc(axis=1, regex=True)["B"]
+
+The regex matching will only work when looking up single strings, not list of strings etc.
+
+.. ipython:: python
+
+    df.loc(regex=True)[["A"], "A"]
+
+*Notice*: Is is currently not possible to set values for a given regular expression.
+
+.. ipython:: python
+
+    df.loc(regex=True)["B", "B"] = [[1, 2], [3, 4]]
+
+
 .. _indexing.deprecate_ix:
 
 IX indexer is deprecated

doc/source/whatsnew/v0.25.0.rst

@@ -182,6 +182,35 @@ The repr now looks like this:
     json_normalize(data, max_level=1)
 
 
+.. _whatsnew_0250.enhancements.loc_regex:
+
+Selection using regular expressions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+It is now possible to call :attr:`~DataFrame.loc` with parameter ``regex=True`` to select by
+row/columns axis labels that match a regular expression pattern.
+
+.. ipython:: python
+
+    df = pd.DataFrame(1, index=["A", "AB", "BC", "CD"], columns=["CD", "BC", "AB", "A"])
+    df
+
+    df.loc(regex=True)["B", "B"]
+    df.loc(axis=1, regex=True)["B"]
+
+The regex matching will only work when looking up single strings, not list of strings etc.
+
+.. ipython:: python
+
+    df.loc(regex=True)[["A"], "A"]
+
+Is is currently not possible to set values for a given regular expression.
+
+.. ipython:: python
+
+    df.loc(regex=True)["B", "B"] = [[1, 2], [3, 4]]
+
+
 .. _whatsnew_0250.enhancements.other:
 
 Other enhancements

pandas/core/indexing.py

@@ -942,11 +942,15 @@ def _getitem_lowerdim(self, tup):
                 # slice returns a new object.
                 if com.is_null_slice(new_key):
                     return section
-                # This is an elided recursive call to iloc/loc/etc'
-                return getattr(section, self.name)[new_key]
+
+                return self._getitem_lower_dim(section, new_key)
 
         raise IndexingError("not applicable")
 
+    def _getitem_lower_dim(self, section, key):
+        # This is an elided recursive call to iloc/loc/etc'
+        return getattr(section, self.name)[key]
+
     def _getitem_nested_tuple(self, tup):
         # we have a nested tuple so have at least 1 multi-index level
         # we should be able to match up the dimensionality here
@@ -1692,6 +1696,14 @@ class _LocIndexer(_LocationIndexer):
     )
     _exception = KeyError
 
+    regex = False
+
+    def __call__(self, axis=None, regex=False):
+        new_self = super().__call__(axis=axis)
+        if regex:
+            new_self.regex = regex
+        return new_self
+
     @Appender(_NDFrameIndexer._validate_key.__doc__)
     def _validate_key(self, key, axis: int):
 
@@ -1761,6 +1773,27 @@ def _get_partial_string_timestamp_match_key(self, key, labels):
 
         return key
 
+    def _get_regex_mappings(self, key, axis=None):
+        import re
+
+        if axis is None:
+            axis = self.axis or 0
+
+        labels = self.obj._get_axis(axis)
+
+        matcher = re.compile(key)
+
+        def f(x):
+            return matcher.search(x) is not None
+
+        return labels.map(f)
+
+    def _getitem_regex(self, key, axis=None):
+        """Subset obj by regex-searching axis for key."""
+        mapped = self._get_regex_mappings(key, axis)
+
+        return self.obj.loc(axis=axis)[mapped]
+
     def _getitem_axis(self, key, axis: int):
         key = item_from_zerodim(key)
         if is_iterator(key):
@@ -1822,10 +1855,21 @@ def _getitem_axis(self, key, axis: int):
                 indexer[axis] = locs
                 return self.obj.iloc[tuple(indexer)]
 
+        if self.regex and isinstance(key, str):
+            return self._getitem_regex(key, axis=axis)
+
         # fall thru to straight lookup
         self._validate_key(key, axis)
         return self._get_label(key, axis=axis)
 
+    def _getitem_lower_dim(self, section, key):
+        return getattr(section, self.name)(regex=self.regex)[key]
+
+    def __setitem__(self, key, value):
+        if self.regex:
+            raise TypeError("Inserting with regex not supported")
+        return super().__setitem__(key, value)
+
 
 class _iLocIndexer(_LocationIndexer):
     """

pandas/tests/indexing/test_loc.py

@@ -13,6 +13,66 @@
 from pandas.util import testing as tm
 
 
+class TestLocRegex:
+    # test calls to df.loc, with a regex parameter set to True,
+    # ie. df.loc(..., regex=True)[...]
+
+    idx = ["AB", "BC", "CD", "DE"]
+    cols = idx[::-1]
+
+    def test_regex_frame(self):
+        idx, cols = self.idx, self.cols
+
+        df = pd.DataFrame(1, index=idx, columns=cols)
+        ser = df["AB"]
+
+        result = df.loc(regex=True)["B"]
+        expected = pd.DataFrame(1, index=["AB", "BC"], columns=cols)
+        tm.assert_frame_equal(result, expected)
+
+        result = ser.loc(regex=True)["B"]
+        expected = pd.Series(1, index=["AB", "BC"], name="AB")
+        tm.assert_series_equal(result, expected)
+
+        result = df.loc(regex=True)[:, "B"]
+        expected = pd.DataFrame(1, index=idx, columns=["BC", "AB"])
+        tm.assert_frame_equal(result, expected)
+
+        result = df.loc(regex=True)["B", "B"]
+        expected = pd.DataFrame(1, index=["AB", "BC"], columns=["BC", "AB"])
+        tm.assert_frame_equal(result, expected)
+
+    def test_regex_empty(self):
+        idx, cols = self.idx, self.cols
+
+        df = pd.DataFrame(1, index=idx, columns=cols)
+        ser = df["AB"]
+
+        result = df.loc(regex=True)["X"]
+        expected = pd.DataFrame(columns=cols, dtype="int64")
+        tm.assert_frame_equal(result, expected)
+
+        result = ser.loc(regex=True)["X"]
+        expected = pd.Series(name="AB", dtype="int64")
+        tm.assert_series_equal(result, expected)
+
+        result = df.loc(regex=True)[:, "X"]
+        expected = pd.DataFrame(index=idx, dtype="int64")
+        tm.assert_frame_equal(result, expected)
+
+        result = df.loc(regex=True)["X", "X"]
+        expected = pd.DataFrame(dtype="int64")
+        tm.assert_frame_equal(result, expected)
+
+    def test_regex_inserting(self):
+        idx, cols = self.idx, self.cols
+
+        df = pd.DataFrame(1, index=idx, columns=cols)
+
+        with pytest.raises(TypeError, match="Inserting with regex not supported"):
+            df.loc(regex=True)["B", "B"] = [[2, 2], [2, 2]]
+
+
 class TestLoc(Base):
     def test_loc_getitem_dups(self):
         # GH 5678