diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 5077b17ba1e..abba8167b65 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -182,6 +182,10 @@ Deprecations Bug Fixes ~~~~~~~~~ +- ``DataArray.str.extract`` with a named output dimension now supports NumPy + variable-width strings, including unmatched values and optional capture groups + (:issue:`11785`). + By `jdotc1 `_. - :py:meth:`Dataset.query` now resolves coordinates in datasets with no data variables (:issue:`6449`). By `Liang, Jingyi `_. diff --git a/xarray/core/accessor_str.py b/xarray/core/accessor_str.py index a51c3822bc2..5be9bff2dcb 100644 --- a/xarray/core/accessor_str.py +++ b/xarray/core/accessor_str.py @@ -2154,10 +2154,10 @@ def _get_res_single(val, pat): def _get_res_multi(val, pat): match = pat.search(val) if match is None: - return np.array([""], val.dtype) + return np.array([""], dtype=np.object_) match = match.groups() match = [grp if grp is not None else "" for grp in match] - return np.array(match, val.dtype) + return np.array(match, dtype=np.object_) if dim is None: return self._apply(func=_get_res_single, func_args=(pat,)) diff --git a/xarray/tests/test_accessor_str.py b/xarray/tests/test_accessor_str.py index d49f8bca2c9..e89db57794c 100644 --- a/xarray/tests/test_accessor_str.py +++ b/xarray/tests/test_accessor_str.py @@ -852,6 +852,29 @@ def test_extract_single_nocase(dtype) -> None: assert_equal(res_re_dim, targ_dim) +@pytest.mark.parametrize("string_dtype", ["T", "U", "S"]) +@pytest.mark.parametrize("compiled", [False, True]) +def test_extract_multi_string_dtypes(string_dtype, compiled) -> None: + value = xr.DataArray( + np.array(["alpha_1234", "beta", "!"], dtype=string_dtype), dims="X" + ) + pattern_str = r"([a-z]+)(?:_(\d+))?" + pattern = ( + re.compile(pattern_str.encode() if string_dtype == "S" else pattern_str) + if compiled + else pattern_str + ) + expected = xr.DataArray( + np.array([["alpha", "1234"], ["beta", ""], ["", ""]], dtype=string_dtype), + dims=["X", "group"], + ) + + result = value.str.extract(pattern, dim="group") + + assert result.dtype == expected.dtype + assert_equal(result, expected) + + def test_extract_multi_case(dtype) -> None: pat_str = r"(\w+)_Xy_(\d*)" pat_re: str | bytes = (