diff --git a/doc/whats-new.rst b/doc/whats-new.rst index d1505bfa081..8c807f03dab 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -53,6 +53,11 @@ Deprecations Bug Fixes ~~~~~~~~~ +- Fixed ``DataArray.str.replace`` replacing every occurrence instead of none when + ``n=0``. ``re.sub`` treats ``count=0`` as "replace all", so the regex code path + collapsed ``n=0`` onto ``n=-1``, while the ``regex=False`` path already handled + ``n=0`` correctly (:pull:`11545`). + By `Alexander Kropiunig `_. - Fix async zarr tests using ``wraps`` with ``autospec=True`` on async methods, which caused ``AsyncMock`` objects to leak through instead of real array data (:pull:`11232`). diff --git a/xarray/core/accessor_str.py b/xarray/core/accessor_str.py index 0699fbdd5b1..b2f6fd4a897 100644 --- a/xarray/core/accessor_str.py +++ b/xarray/core/accessor_str.py @@ -1944,8 +1944,10 @@ def replace( if regex: pat = self._re_compile(pat=pat, flags=flags, case=case) - func = lambda x, ipat, irepl, i_n: ipat.sub( - repl=irepl, string=x, count=max(i_n, 0) + # ``re.sub`` interprets ``count=0`` as "replace every occurrence", + # so ``n=0`` has to be special-cased to mean "replace nothing". + func = lambda x, ipat, irepl, i_n: ( + x if i_n == 0 else ipat.sub(repl=irepl, string=x, count=max(i_n, 0)) ) else: pat = self._stringify(pat) diff --git a/xarray/tests/test_accessor_str.py b/xarray/tests/test_accessor_str.py index 0741fa364fc..5f5849bc9de 100644 --- a/xarray/tests/test_accessor_str.py +++ b/xarray/tests/test_accessor_str.py @@ -381,6 +381,28 @@ def test_replace(dtype) -> None: assert_equal(result, expected) +def test_replace_n_zero(dtype) -> None: + # ``n=0`` means "make no replacements", for regex and literal patterns alike + values = xr.DataArray(["fooBAD__barBAD"], dims=["x"]).astype(dtype) + + result = values.str.replace("BAD[_]*", "", n=0) + assert result.dtype == values.dtype + assert_equal(result, values) + + result = values.str.replace("BAD", "", n=0, regex=False) + assert result.dtype == values.dtype + assert_equal(result, values) + + # ``n`` is broadcast, so a single zero must not spill over to its neighbours + n = xr.DataArray([0, 1, -1], dims=["y"]) + result = values.str.replace("BAD[_]*", "", n=n) + expected = xr.DataArray( + [["fooBAD__barBAD", "foobarBAD", "foobar"]], dims=["x", "y"] + ).astype(dtype) + assert result.dtype == expected.dtype + assert_equal(result, expected) + + def test_replace_callable() -> None: values = xr.DataArray(["fooBAD__barBAD"])