File: C:/Users/fred/anaconda3/Lib/site-packages/xarray/tests/test_accessor_str.py
# Tests for the `str` accessor are derived from the original
# pandas string accessor tests.
# For reference, here is a copy of the pandas copyright notice:
# (c) 2011-2012, Lambda Foundry, Inc. and PyData Development Team
# All rights reserved.
# Copyright (c) 2008-2011 AQR Capital Management, LLC
# All rights reserved.
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following
# disclaimer in the documentation and/or other materials provided
# with the distribution.
# * Neither the name of the copyright holder nor the names of any
# contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from __future__ import annotations
import re
from typing import Callable
import numpy as np
import pytest
import xarray as xr
from xarray.tests import assert_equal, assert_identical, requires_dask
@pytest.fixture(
params=[pytest.param(np.str_, id="str"), pytest.param(np.bytes_, id="bytes")]
)
def dtype(request):
return request.param
@requires_dask
def test_dask() -> None:
import dask.array as da
arr = da.from_array(["a", "b", "c"], chunks=-1)
xarr = xr.DataArray(arr)
result = xarr.str.len().compute()
expected = xr.DataArray([1, 1, 1])
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_count(dtype) -> None:
values = xr.DataArray(["foo", "foofoo", "foooofooofommmfoo"]).astype(dtype)
pat_str = dtype(r"f[o]+")
pat_re = re.compile(pat_str)
result_str = values.str.count(pat_str)
result_re = values.str.count(pat_re)
expected = xr.DataArray([1, 2, 4])
assert result_str.dtype == expected.dtype
assert result_re.dtype == expected.dtype
assert_equal(result_str, expected)
assert_equal(result_re, expected)
def test_count_broadcast(dtype) -> None:
values = xr.DataArray(["foo", "foofoo", "foooofooofommmfoo"]).astype(dtype)
pat_str = np.array([r"f[o]+", r"o", r"m"]).astype(dtype)
pat_re = np.array([re.compile(x) for x in pat_str])
result_str = values.str.count(pat_str)
result_re = values.str.count(pat_re)
expected = xr.DataArray([1, 4, 3])
assert result_str.dtype == expected.dtype
assert result_re.dtype == expected.dtype
assert_equal(result_str, expected)
assert_equal(result_re, expected)
def test_contains(dtype) -> None:
values = xr.DataArray(["Foo", "xYz", "fOOomMm__fOo", "MMM_"]).astype(dtype)
# case insensitive using regex
pat = values.dtype.type("FOO|mmm")
result = values.str.contains(pat, case=False)
expected = xr.DataArray([True, False, True, True])
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.contains(re.compile(pat, flags=re.IGNORECASE))
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case sensitive using regex
pat = values.dtype.type("Foo|mMm")
result = values.str.contains(pat)
expected = xr.DataArray([True, False, True, False])
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.contains(re.compile(pat))
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case insensitive without regex
result = values.str.contains("foo", regex=False, case=False)
expected = xr.DataArray([True, False, True, False])
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case sensitive without regex
result = values.str.contains("fO", regex=False, case=True)
expected = xr.DataArray([False, False, True, False])
assert result.dtype == expected.dtype
assert_equal(result, expected)
# regex regex=False
pat_re = re.compile("(/w+)")
with pytest.raises(
ValueError,
match="Must use regular expression matching for regular expression object.",
):
values.str.contains(pat_re, regex=False)
def test_contains_broadcast(dtype) -> None:
values = xr.DataArray(["Foo", "xYz", "fOOomMm__fOo", "MMM_"], dims="X").astype(
dtype
)
pat_str = xr.DataArray(["FOO|mmm", "Foo", "MMM"], dims="Y").astype(dtype)
pat_re = xr.DataArray([re.compile(x) for x in pat_str.data], dims="Y")
# case insensitive using regex
result = values.str.contains(pat_str, case=False)
expected = xr.DataArray(
[
[True, True, False],
[False, False, False],
[True, True, True],
[True, False, True],
],
dims=["X", "Y"],
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case sensitive using regex
result = values.str.contains(pat_str)
expected = xr.DataArray(
[
[False, True, False],
[False, False, False],
[False, False, False],
[False, False, True],
],
dims=["X", "Y"],
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.contains(pat_re)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case insensitive without regex
result = values.str.contains(pat_str, regex=False, case=False)
expected = xr.DataArray(
[
[False, True, False],
[False, False, False],
[False, True, True],
[False, False, True],
],
dims=["X", "Y"],
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case insensitive with regex
result = values.str.contains(pat_str, regex=False, case=True)
expected = xr.DataArray(
[
[False, True, False],
[False, False, False],
[False, False, False],
[False, False, True],
],
dims=["X", "Y"],
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_starts_ends_with(dtype) -> None:
values = xr.DataArray(["om", "foo_nom", "nom", "bar_foo", "foo"]).astype(dtype)
result = values.str.startswith("foo")
expected = xr.DataArray([False, True, False, False, True])
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.endswith("foo")
expected = xr.DataArray([False, False, False, True, True])
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_starts_ends_with_broadcast(dtype) -> None:
values = xr.DataArray(
["om", "foo_nom", "nom", "bar_foo", "foo_bar"], dims="X"
).astype(dtype)
pat = xr.DataArray(["foo", "bar"], dims="Y").astype(dtype)
result = values.str.startswith(pat)
expected = xr.DataArray(
[[False, False], [True, False], [False, False], [False, True], [True, False]],
dims=["X", "Y"],
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.endswith(pat)
expected = xr.DataArray(
[[False, False], [False, False], [False, False], [True, False], [False, True]],
dims=["X", "Y"],
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_case_bytes() -> None:
value = xr.DataArray(["SOme wOrd"]).astype(np.bytes_)
exp_capitalized = xr.DataArray(["Some word"]).astype(np.bytes_)
exp_lowered = xr.DataArray(["some word"]).astype(np.bytes_)
exp_swapped = xr.DataArray(["soME WoRD"]).astype(np.bytes_)
exp_titled = xr.DataArray(["Some Word"]).astype(np.bytes_)
exp_uppered = xr.DataArray(["SOME WORD"]).astype(np.bytes_)
res_capitalized = value.str.capitalize()
res_lowered = value.str.lower()
res_swapped = value.str.swapcase()
res_titled = value.str.title()
res_uppered = value.str.upper()
assert res_capitalized.dtype == exp_capitalized.dtype
assert res_lowered.dtype == exp_lowered.dtype
assert res_swapped.dtype == exp_swapped.dtype
assert res_titled.dtype == exp_titled.dtype
assert res_uppered.dtype == exp_uppered.dtype
assert_equal(res_capitalized, exp_capitalized)
assert_equal(res_lowered, exp_lowered)
assert_equal(res_swapped, exp_swapped)
assert_equal(res_titled, exp_titled)
assert_equal(res_uppered, exp_uppered)
def test_case_str() -> None:
# This string includes some unicode characters
# that are common case management corner cases
value = xr.DataArray(["SOme wOrd DŽ ß ᾛ ΣΣ ffi⁵Å Ç Ⅰ"]).astype(np.unicode_)
exp_capitalized = xr.DataArray(["Some word dž ß ᾓ σς ffi⁵å ç ⅰ"]).astype(np.unicode_)
exp_lowered = xr.DataArray(["some word dž ß ᾓ σς ffi⁵å ç ⅰ"]).astype(np.unicode_)
exp_swapped = xr.DataArray(["soME WoRD dž SS ᾛ σς FFI⁵å ç ⅰ"]).astype(np.unicode_)
exp_titled = xr.DataArray(["Some Word Dž Ss ᾛ Σς Ffi⁵Å Ç Ⅰ"]).astype(np.unicode_)
exp_uppered = xr.DataArray(["SOME WORD DŽ SS ἫΙ ΣΣ FFI⁵Å Ç Ⅰ"]).astype(np.unicode_)
exp_casefolded = xr.DataArray(["some word dž ss ἣι σσ ffi⁵å ç ⅰ"]).astype(
np.unicode_
)
exp_norm_nfc = xr.DataArray(["SOme wOrd DŽ ß ᾛ ΣΣ ffi⁵Å Ç Ⅰ"]).astype(np.unicode_)
exp_norm_nfkc = xr.DataArray(["SOme wOrd DŽ ß ᾛ ΣΣ ffi5Å Ç I"]).astype(np.unicode_)
exp_norm_nfd = xr.DataArray(["SOme wOrd DŽ ß ᾛ ΣΣ ffi⁵Å Ç Ⅰ"]).astype(np.unicode_)
exp_norm_nfkd = xr.DataArray(["SOme wOrd DŽ ß ᾛ ΣΣ ffi5Å Ç I"]).astype(
np.unicode_
)
res_capitalized = value.str.capitalize()
res_casefolded = value.str.casefold()
res_lowered = value.str.lower()
res_swapped = value.str.swapcase()
res_titled = value.str.title()
res_uppered = value.str.upper()
res_norm_nfc = value.str.normalize("NFC")
res_norm_nfd = value.str.normalize("NFD")
res_norm_nfkc = value.str.normalize("NFKC")
res_norm_nfkd = value.str.normalize("NFKD")
assert res_capitalized.dtype == exp_capitalized.dtype
assert res_casefolded.dtype == exp_casefolded.dtype
assert res_lowered.dtype == exp_lowered.dtype
assert res_swapped.dtype == exp_swapped.dtype
assert res_titled.dtype == exp_titled.dtype
assert res_uppered.dtype == exp_uppered.dtype
assert res_norm_nfc.dtype == exp_norm_nfc.dtype
assert res_norm_nfd.dtype == exp_norm_nfd.dtype
assert res_norm_nfkc.dtype == exp_norm_nfkc.dtype
assert res_norm_nfkd.dtype == exp_norm_nfkd.dtype
assert_equal(res_capitalized, exp_capitalized)
assert_equal(res_casefolded, exp_casefolded)
assert_equal(res_lowered, exp_lowered)
assert_equal(res_swapped, exp_swapped)
assert_equal(res_titled, exp_titled)
assert_equal(res_uppered, exp_uppered)
assert_equal(res_norm_nfc, exp_norm_nfc)
assert_equal(res_norm_nfd, exp_norm_nfd)
assert_equal(res_norm_nfkc, exp_norm_nfkc)
assert_equal(res_norm_nfkd, exp_norm_nfkd)
def test_replace(dtype) -> None:
values = xr.DataArray(["fooBAD__barBAD"], dims=["x"]).astype(dtype)
result = values.str.replace("BAD[_]*", "")
expected = xr.DataArray(["foobar"], dims=["x"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.replace("BAD[_]*", "", n=1)
expected = xr.DataArray(["foobarBAD"], dims=["x"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
pat = xr.DataArray(["BAD[_]*", "AD[_]*"], dims=["y"]).astype(dtype)
result = values.str.replace(pat, "")
expected = xr.DataArray([["foobar", "fooBbarB"]], dims=["x", "y"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
repl = xr.DataArray(["", "spam"], dims=["y"]).astype(dtype)
result = values.str.replace(pat, repl, n=1)
expected = xr.DataArray([["foobarBAD", "fooBspambarBAD"]], dims=["x", "y"]).astype(
dtype
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
values = xr.DataArray(
["A", "B", "C", "Aaba", "Baca", "", "CABA", "dog", "cat"]
).astype(dtype)
expected = xr.DataArray(
["YYY", "B", "C", "YYYaba", "Baca", "", "CYYYBYYY", "dog", "cat"]
).astype(dtype)
result = values.str.replace("A", "YYY")
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.replace("A", "YYY", regex=False)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.replace("A", "YYY", case=False)
expected = xr.DataArray(
["YYY", "B", "C", "YYYYYYbYYY", "BYYYcYYY", "", "CYYYBYYY", "dog", "cYYYt"]
).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.replace("^.a|dog", "XX-XX ", case=False)
expected = xr.DataArray(
["A", "B", "C", "XX-XX ba", "XX-XX ca", "", "XX-XX BA", "XX-XX ", "XX-XX t"]
).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_replace_callable() -> None:
values = xr.DataArray(["fooBAD__barBAD"])
# test with callable
repl = lambda m: m.group(0).swapcase()
result = values.str.replace("[a-z][A-Z]{2}", repl, n=2)
exp = xr.DataArray(["foObaD__baRbaD"])
assert result.dtype == exp.dtype
assert_equal(result, exp)
# test regex named groups
values = xr.DataArray(["Foo Bar Baz"])
pat = r"(?P<first>\w+) (?P<middle>\w+) (?P<last>\w+)"
repl = lambda m: m.group("middle").swapcase()
result = values.str.replace(pat, repl)
exp = xr.DataArray(["bAR"])
assert result.dtype == exp.dtype
assert_equal(result, exp)
# test broadcast
values = xr.DataArray(["Foo Bar Baz"], dims=["x"])
pat = r"(?P<first>\w+) (?P<middle>\w+) (?P<last>\w+)"
repl2 = xr.DataArray(
[
lambda m: m.group("first").swapcase(),
lambda m: m.group("middle").swapcase(),
lambda m: m.group("last").swapcase(),
],
dims=["Y"],
)
result = values.str.replace(pat, repl2)
exp = xr.DataArray([["fOO", "bAR", "bAZ"]], dims=["x", "Y"])
assert result.dtype == exp.dtype
assert_equal(result, exp)
def test_replace_unicode() -> None:
# flags + unicode
values = xr.DataArray([b"abcd,\xc3\xa0".decode("utf-8")])
expected = xr.DataArray([b"abcd, \xc3\xa0".decode("utf-8")])
pat = re.compile(r"(?<=\w),(?=\w)", flags=re.UNICODE)
result = values.str.replace(pat, ", ")
assert result.dtype == expected.dtype
assert_equal(result, expected)
# broadcast version
values = xr.DataArray([b"abcd,\xc3\xa0".decode("utf-8")], dims=["X"])
expected = xr.DataArray(
[[b"abcd, \xc3\xa0".decode("utf-8"), b"BAcd,\xc3\xa0".decode("utf-8")]],
dims=["X", "Y"],
)
pat2 = xr.DataArray(
[re.compile(r"(?<=\w),(?=\w)", flags=re.UNICODE), r"ab"], dims=["Y"]
)
repl = xr.DataArray([", ", "BA"], dims=["Y"])
result = values.str.replace(pat2, repl)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_replace_compiled_regex(dtype) -> None:
values = xr.DataArray(["fooBAD__barBAD"], dims=["x"]).astype(dtype)
# test with compiled regex
pat = re.compile(dtype("BAD[_]*"))
result = values.str.replace(pat, "")
expected = xr.DataArray(["foobar"], dims=["x"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.replace(pat, "", n=1)
expected = xr.DataArray(["foobarBAD"], dims=["x"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# broadcast
pat2 = xr.DataArray(
[re.compile(dtype("BAD[_]*")), re.compile(dtype("AD[_]*"))], dims=["y"]
)
result = values.str.replace(pat2, "")
expected = xr.DataArray([["foobar", "fooBbarB"]], dims=["x", "y"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
repl = xr.DataArray(["", "spam"], dims=["y"]).astype(dtype)
result = values.str.replace(pat2, repl, n=1)
expected = xr.DataArray([["foobarBAD", "fooBspambarBAD"]], dims=["x", "y"]).astype(
dtype
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# case and flags provided to str.replace will have no effect
# and will produce warnings
values = xr.DataArray(["fooBAD__barBAD__bad"]).astype(dtype)
pat3 = re.compile(dtype("BAD[_]*"))
with pytest.raises(
ValueError, match="Flags cannot be set when pat is a compiled regex."
):
result = values.str.replace(pat3, "", flags=re.IGNORECASE)
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
result = values.str.replace(pat3, "", case=False)
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
result = values.str.replace(pat3, "", case=True)
# test with callable
values = xr.DataArray(["fooBAD__barBAD"]).astype(dtype)
repl2 = lambda m: m.group(0).swapcase()
pat4 = re.compile(dtype("[a-z][A-Z]{2}"))
result = values.str.replace(pat4, repl2, n=2)
expected = xr.DataArray(["foObaD__baRbaD"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_replace_literal(dtype) -> None:
# GH16808 literal replace (regex=False vs regex=True)
values = xr.DataArray(["f.o", "foo"], dims=["X"]).astype(dtype)
expected = xr.DataArray(["bao", "bao"], dims=["X"]).astype(dtype)
result = values.str.replace("f.", "ba")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = xr.DataArray(["bao", "foo"], dims=["X"]).astype(dtype)
result = values.str.replace("f.", "ba", regex=False)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# Broadcast
pat = xr.DataArray(["f.", ".o"], dims=["yy"]).astype(dtype)
expected = xr.DataArray([["bao", "fba"], ["bao", "bao"]], dims=["X", "yy"]).astype(
dtype
)
result = values.str.replace(pat, "ba")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = xr.DataArray([["bao", "fba"], ["foo", "foo"]], dims=["X", "yy"]).astype(
dtype
)
result = values.str.replace(pat, "ba", regex=False)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# Cannot do a literal replace if given a callable repl or compiled
# pattern
callable_repl = lambda m: m.group(0).swapcase()
compiled_pat = re.compile("[a-z][A-Z]{2}")
msg = "Cannot use a callable replacement when regex=False"
with pytest.raises(ValueError, match=msg):
values.str.replace("abc", callable_repl, regex=False)
msg = "Cannot use a compiled regex as replacement pattern with regex=False"
with pytest.raises(ValueError, match=msg):
values.str.replace(compiled_pat, "", regex=False)
def test_extract_extractall_findall_empty_raises(dtype) -> None:
pat_str = dtype(r".*")
pat_re = re.compile(pat_str)
value = xr.DataArray([["a"]], dims=["X", "Y"]).astype(dtype)
with pytest.raises(ValueError, match="No capture groups found in pattern."):
value.str.extract(pat=pat_str, dim="ZZ")
with pytest.raises(ValueError, match="No capture groups found in pattern."):
value.str.extract(pat=pat_re, dim="ZZ")
with pytest.raises(ValueError, match="No capture groups found in pattern."):
value.str.extractall(pat=pat_str, group_dim="XX", match_dim="YY")
with pytest.raises(ValueError, match="No capture groups found in pattern."):
value.str.extractall(pat=pat_re, group_dim="XX", match_dim="YY")
with pytest.raises(ValueError, match="No capture groups found in pattern."):
value.str.findall(pat=pat_str)
with pytest.raises(ValueError, match="No capture groups found in pattern."):
value.str.findall(pat=pat_re)
def test_extract_multi_None_raises(dtype) -> None:
pat_str = r"(\w+)_(\d+)"
pat_re = re.compile(pat_str)
value = xr.DataArray([["a_b"]], dims=["X", "Y"]).astype(dtype)
with pytest.raises(
ValueError,
match="Dimension must be specified if more than one capture group is given.",
):
value.str.extract(pat=pat_str, dim=None)
with pytest.raises(
ValueError,
match="Dimension must be specified if more than one capture group is given.",
):
value.str.extract(pat=pat_re, dim=None)
def test_extract_extractall_findall_case_re_raises(dtype) -> None:
pat_str = r".*"
pat_re = re.compile(pat_str)
value = xr.DataArray([["a"]], dims=["X", "Y"]).astype(dtype)
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
value.str.extract(pat=pat_re, case=True, dim="ZZ")
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
value.str.extract(pat=pat_re, case=False, dim="ZZ")
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
value.str.extractall(pat=pat_re, case=True, group_dim="XX", match_dim="YY")
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
value.str.extractall(pat=pat_re, case=False, group_dim="XX", match_dim="YY")
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
value.str.findall(pat=pat_re, case=True)
with pytest.raises(
ValueError, match="Case cannot be set when pat is a compiled regex."
):
value.str.findall(pat=pat_re, case=False)
def test_extract_extractall_name_collision_raises(dtype) -> None:
pat_str = r"(\w+)"
pat_re = re.compile(pat_str)
value = xr.DataArray([["a"]], dims=["X", "Y"]).astype(dtype)
with pytest.raises(KeyError, match="Dimension 'X' already present in DataArray."):
value.str.extract(pat=pat_str, dim="X")
with pytest.raises(KeyError, match="Dimension 'X' already present in DataArray."):
value.str.extract(pat=pat_re, dim="X")
with pytest.raises(
KeyError, match="Group dimension 'X' already present in DataArray."
):
value.str.extractall(pat=pat_str, group_dim="X", match_dim="ZZ")
with pytest.raises(
KeyError, match="Group dimension 'X' already present in DataArray."
):
value.str.extractall(pat=pat_re, group_dim="X", match_dim="YY")
with pytest.raises(
KeyError, match="Match dimension 'Y' already present in DataArray."
):
value.str.extractall(pat=pat_str, group_dim="XX", match_dim="Y")
with pytest.raises(
KeyError, match="Match dimension 'Y' already present in DataArray."
):
value.str.extractall(pat=pat_re, group_dim="XX", match_dim="Y")
with pytest.raises(
KeyError, match="Group dimension 'ZZ' is the same as match dimension 'ZZ'."
):
value.str.extractall(pat=pat_str, group_dim="ZZ", match_dim="ZZ")
with pytest.raises(
KeyError, match="Group dimension 'ZZ' is the same as match dimension 'ZZ'."
):
value.str.extractall(pat=pat_re, group_dim="ZZ", match_dim="ZZ")
def test_extract_single_case(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
targ_none = xr.DataArray(
[["a", "bab", "abc"], ["abcd", "", "abcdef"]], dims=["X", "Y"]
).astype(dtype)
targ_dim = xr.DataArray(
[[["a"], ["bab"], ["abc"]], [["abcd"], [""], ["abcdef"]]], dims=["X", "Y", "XX"]
).astype(dtype)
res_str_none = value.str.extract(pat=pat_str, dim=None)
res_str_dim = value.str.extract(pat=pat_str, dim="XX")
res_str_none_case = value.str.extract(pat=pat_str, dim=None, case=True)
res_str_dim_case = value.str.extract(pat=pat_str, dim="XX", case=True)
res_re_none = value.str.extract(pat=pat_compiled, dim=None)
res_re_dim = value.str.extract(pat=pat_compiled, dim="XX")
assert res_str_none.dtype == targ_none.dtype
assert res_str_dim.dtype == targ_dim.dtype
assert res_str_none_case.dtype == targ_none.dtype
assert res_str_dim_case.dtype == targ_dim.dtype
assert res_re_none.dtype == targ_none.dtype
assert res_re_dim.dtype == targ_dim.dtype
assert_equal(res_str_none, targ_none)
assert_equal(res_str_dim, targ_dim)
assert_equal(res_str_none_case, targ_none)
assert_equal(res_str_dim_case, targ_dim)
assert_equal(res_re_none, targ_none)
assert_equal(res_re_dim, targ_dim)
def test_extract_single_nocase(dtype) -> None:
pat_str = r"(\w+)?_Xy_\d*"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re, flags=re.IGNORECASE)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"_Xy_1",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
targ_none = xr.DataArray(
[["a", "ab", "abc"], ["abcd", "", "abcdef"]], dims=["X", "Y"]
).astype(dtype)
targ_dim = xr.DataArray(
[[["a"], ["ab"], ["abc"]], [["abcd"], [""], ["abcdef"]]], dims=["X", "Y", "XX"]
).astype(dtype)
res_str_none = value.str.extract(pat=pat_str, dim=None, case=False)
res_str_dim = value.str.extract(pat=pat_str, dim="XX", case=False)
res_re_none = value.str.extract(pat=pat_compiled, dim=None)
res_re_dim = value.str.extract(pat=pat_compiled, dim="XX")
assert res_re_dim.dtype == targ_none.dtype
assert res_str_dim.dtype == targ_dim.dtype
assert res_re_none.dtype == targ_none.dtype
assert res_re_dim.dtype == targ_dim.dtype
assert_equal(res_str_none, targ_none)
assert_equal(res_str_dim, targ_dim)
assert_equal(res_re_none, targ_none)
assert_equal(res_re_dim, targ_dim)
def test_extract_multi_case(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[["a", "0"], ["bab", "110"], ["abc", "01"]],
[["abcd", ""], ["", ""], ["abcdef", "101"]],
],
dims=["X", "Y", "XX"],
).astype(dtype)
res_str = value.str.extract(pat=pat_str, dim="XX")
res_re = value.str.extract(pat=pat_compiled, dim="XX")
res_str_case = value.str.extract(pat=pat_str, dim="XX", case=True)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_extract_multi_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re, flags=re.IGNORECASE)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[["a", "0"], ["ab", "10"], ["abc", "01"]],
[["abcd", ""], ["", ""], ["abcdef", "101"]],
],
dims=["X", "Y", "XX"],
).astype(dtype)
res_str = value.str.extract(pat=pat_str, dim="XX", case=False)
res_re = value.str.extract(pat=pat_compiled, dim="XX")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_extract_broadcast(dtype) -> None:
value = xr.DataArray(
["a_Xy_0", "ab_xY_10", "abc_Xy_01"],
dims=["X"],
).astype(dtype)
pat_str = xr.DataArray(
[r"(\w+)_Xy_(\d*)", r"(\w+)_xY_(\d*)"],
dims=["Y"],
).astype(dtype)
pat_compiled = value.str._re_compile(pat=pat_str)
expected_list = [
[["a", "0"], ["", ""]],
[["", ""], ["ab", "10"]],
[["abc", "01"], ["", ""]],
]
expected = xr.DataArray(expected_list, dims=["X", "Y", "Zz"]).astype(dtype)
res_str = value.str.extract(pat=pat_str, dim="Zz")
res_re = value.str.extract(pat=pat_compiled, dim="Zz")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_extractall_single_single_case(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re)
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[[[["a"]], [[""]], [["abc"]]], [[["abcd"]], [[""]], [["abcdef"]]]],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(pat=pat_str, group_dim="XX", match_dim="YY")
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
res_str_case = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=True
)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_extractall_single_single_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re, flags=re.I)
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[[[["a"]], [["ab"]], [["abc"]]], [[["abcd"]], [[""]], [["abcdef"]]]],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=False
)
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_extractall_single_multi_case(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[[["a"], [""], [""]], [["bab"], ["baab"], [""]], [["abc"], ["cbc"], [""]]],
[
[["abcd"], ["dcd"], ["dccd"]],
[[""], [""], [""]],
[["abcdef"], ["fef"], [""]],
],
],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(pat=pat_str, group_dim="XX", match_dim="YY")
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
res_str_case = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=True
)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_extractall_single_multi_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re, flags=re.I)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[
[["a"], [""], [""]],
[["ab"], ["bab"], ["baab"]],
[["abc"], ["cbc"], [""]],
],
[
[["abcd"], ["dcd"], ["dccd"]],
[[""], [""], [""]],
[["abcdef"], ["fef"], [""]],
],
],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=False
)
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_extractall_multi_single_case(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re)
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[[["a", "0"]], [["", ""]], [["abc", "01"]]],
[[["abcd", ""]], [["", ""]], [["abcdef", "101"]]],
],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(pat=pat_str, group_dim="XX", match_dim="YY")
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
res_str_case = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=True
)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_extractall_multi_single_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re, flags=re.I)
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[[["a", "0"]], [["ab", "10"]], [["abc", "01"]]],
[[["abcd", ""]], [["", ""]], [["abcdef", "101"]]],
],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=False
)
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_extractall_multi_multi_case(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[
[["a", "0"], ["", ""], ["", ""]],
[["bab", "110"], ["baab", "1100"], ["", ""]],
[["abc", "01"], ["cbc", "2210"], ["", ""]],
],
[
[["abcd", ""], ["dcd", "33210"], ["dccd", "332210"]],
[["", ""], ["", ""], ["", ""]],
[["abcdef", "101"], ["fef", "5543210"], ["", ""]],
],
],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(pat=pat_str, group_dim="XX", match_dim="YY")
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
res_str_case = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=True
)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_extractall_multi_multi_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re: str | bytes = (
pat_str if dtype == np.unicode_ else bytes(pat_str, encoding="UTF-8")
)
pat_compiled = re.compile(pat_re, flags=re.I)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected = xr.DataArray(
[
[
[["a", "0"], ["", ""], ["", ""]],
[["ab", "10"], ["bab", "110"], ["baab", "1100"]],
[["abc", "01"], ["cbc", "2210"], ["", ""]],
],
[
[["abcd", ""], ["dcd", "33210"], ["dccd", "332210"]],
[["", ""], ["", ""], ["", ""]],
[["abcdef", "101"], ["fef", "5543210"], ["", ""]],
],
],
dims=["X", "Y", "XX", "YY"],
).astype(dtype)
res_str = value.str.extractall(
pat=pat_str, group_dim="XX", match_dim="YY", case=False
)
res_re = value.str.extractall(pat=pat_compiled, group_dim="XX", match_dim="YY")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_extractall_broadcast(dtype) -> None:
value = xr.DataArray(
["a_Xy_0", "ab_xY_10", "abc_Xy_01"],
dims=["X"],
).astype(dtype)
pat_str = xr.DataArray(
[r"(\w+)_Xy_(\d*)", r"(\w+)_xY_(\d*)"],
dims=["Y"],
).astype(dtype)
pat_re = value.str._re_compile(pat=pat_str)
expected_list = [
[[["a", "0"]], [["", ""]]],
[[["", ""]], [["ab", "10"]]],
[[["abc", "01"]], [["", ""]]],
]
expected = xr.DataArray(expected_list, dims=["X", "Y", "ZX", "ZY"]).astype(dtype)
res_str = value.str.extractall(pat=pat_str, group_dim="ZX", match_dim="ZY")
res_re = value.str.extractall(pat=pat_re, group_dim="ZX", match_dim="ZY")
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_findall_single_single_case(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re = re.compile(dtype(pat_str))
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list]] = [[["a"], [], ["abc"]], [["abcd"], [], ["abcdef"]]]
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected_list]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str)
res_re = value.str.findall(pat=pat_re)
res_str_case = value.str.findall(pat=pat_str, case=True)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_findall_single_single_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re = re.compile(dtype(pat_str), flags=re.I)
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list]] = [
[["a"], ["ab"], ["abc"]],
[["abcd"], [], ["abcdef"]],
]
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected_list]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str, case=False)
res_re = value.str.findall(pat=pat_re)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_findall_single_multi_case(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re = re.compile(dtype(pat_str))
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list]] = [
[["a"], ["bab", "baab"], ["abc", "cbc"]],
[
["abcd", "dcd", "dccd"],
[],
["abcdef", "fef"],
],
]
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected_list]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str)
res_re = value.str.findall(pat=pat_re)
res_str_case = value.str.findall(pat=pat_str, case=True)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_findall_single_multi_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_\d*"
pat_re = re.compile(dtype(pat_str), flags=re.I)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list]] = [
[
["a"],
["ab", "bab", "baab"],
["abc", "cbc"],
],
[
["abcd", "dcd", "dccd"],
[],
["abcdef", "fef"],
],
]
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected_list]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str, case=False)
res_re = value.str.findall(pat=pat_re)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_findall_multi_single_case(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re = re.compile(dtype(pat_str))
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list[list]]] = [
[[["a", "0"]], [], [["abc", "01"]]],
[[["abcd", ""]], [], [["abcdef", "101"]]],
]
expected_dtype = [
[[tuple(dtype(x) for x in y) for y in z] for z in w] for w in expected_list
]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str)
res_re = value.str.findall(pat=pat_re)
res_str_case = value.str.findall(pat=pat_str, case=True)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_findall_multi_single_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re = re.compile(dtype(pat_str), flags=re.I)
value = xr.DataArray(
[["a_Xy_0", "ab_xY_10", "abc_Xy_01"], ["abcd_Xy_", "", "abcdef_Xy_101"]],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list[list]]] = [
[[["a", "0"]], [["ab", "10"]], [["abc", "01"]]],
[[["abcd", ""]], [], [["abcdef", "101"]]],
]
expected_dtype = [
[[tuple(dtype(x) for x in y) for y in z] for z in w] for w in expected_list
]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str, case=False)
res_re = value.str.findall(pat=pat_re)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_findall_multi_multi_case(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re = re.compile(dtype(pat_str))
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list[list]]] = [
[
[["a", "0"]],
[["bab", "110"], ["baab", "1100"]],
[["abc", "01"], ["cbc", "2210"]],
],
[
[["abcd", ""], ["dcd", "33210"], ["dccd", "332210"]],
[],
[["abcdef", "101"], ["fef", "5543210"]],
],
]
expected_dtype = [
[[tuple(dtype(x) for x in y) for y in z] for z in w] for w in expected_list
]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str)
res_re = value.str.findall(pat=pat_re)
res_str_case = value.str.findall(pat=pat_str, case=True)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert res_str_case.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
assert_equal(res_str_case, expected)
def test_findall_multi_multi_nocase(dtype) -> None:
pat_str = r"(\w+)_Xy_(\d*)"
pat_re = re.compile(dtype(pat_str), flags=re.I)
value = xr.DataArray(
[
["a_Xy_0", "ab_xY_10-bab_Xy_110-baab_Xy_1100", "abc_Xy_01-cbc_Xy_2210"],
[
"abcd_Xy_-dcd_Xy_33210-dccd_Xy_332210",
"",
"abcdef_Xy_101-fef_Xy_5543210",
],
],
dims=["X", "Y"],
).astype(dtype)
expected_list: list[list[list[list]]] = [
[
[["a", "0"]],
[["ab", "10"], ["bab", "110"], ["baab", "1100"]],
[["abc", "01"], ["cbc", "2210"]],
],
[
[["abcd", ""], ["dcd", "33210"], ["dccd", "332210"]],
[],
[["abcdef", "101"], ["fef", "5543210"]],
],
]
expected_dtype = [
[[tuple(dtype(x) for x in y) for y in z] for z in w] for w in expected_list
]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str, case=False)
res_re = value.str.findall(pat=pat_re)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_findall_broadcast(dtype) -> None:
value = xr.DataArray(
["a_Xy_0", "ab_xY_10", "abc_Xy_01"],
dims=["X"],
).astype(dtype)
pat_str = xr.DataArray(
[r"(\w+)_Xy_\d*", r"\w+_Xy_(\d*)"],
dims=["Y"],
).astype(dtype)
pat_re = value.str._re_compile(pat=pat_str)
expected_list: list[list[list]] = [[["a"], ["0"]], [[], []], [["abc"], ["01"]]]
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected_list]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected = xr.DataArray(expected_np, dims=["X", "Y"])
res_str = value.str.findall(pat=pat_str)
res_re = value.str.findall(pat=pat_re)
assert res_str.dtype == expected.dtype
assert res_re.dtype == expected.dtype
assert_equal(res_str, expected)
assert_equal(res_re, expected)
def test_repeat(dtype) -> None:
values = xr.DataArray(["a", "b", "c", "d"]).astype(dtype)
result = values.str.repeat(3)
result_mul = values.str * 3
expected = xr.DataArray(["aaa", "bbb", "ccc", "ddd"]).astype(dtype)
assert result.dtype == expected.dtype
assert result_mul.dtype == expected.dtype
assert_equal(result_mul, expected)
assert_equal(result, expected)
def test_repeat_broadcast(dtype) -> None:
values = xr.DataArray(["a", "b", "c", "d"], dims=["X"]).astype(dtype)
reps = xr.DataArray([3, 4], dims=["Y"])
result = values.str.repeat(reps)
result_mul = values.str * reps
expected = xr.DataArray(
[["aaa", "aaaa"], ["bbb", "bbbb"], ["ccc", "cccc"], ["ddd", "dddd"]],
dims=["X", "Y"],
).astype(dtype)
assert result.dtype == expected.dtype
assert result_mul.dtype == expected.dtype
assert_equal(result_mul, expected)
assert_equal(result, expected)
def test_match(dtype) -> None:
values = xr.DataArray(["fooBAD__barBAD", "foo"]).astype(dtype)
# New match behavior introduced in 0.13
pat = values.dtype.type(".*(BAD[_]+).*(BAD)")
result = values.str.match(pat)
expected = xr.DataArray([True, False])
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.match(re.compile(pat))
assert result.dtype == expected.dtype
assert_equal(result, expected)
# Case-sensitive
pat = values.dtype.type(".*BAD[_]+.*BAD")
result = values.str.match(pat)
expected = xr.DataArray([True, False])
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.match(re.compile(pat))
assert result.dtype == expected.dtype
assert_equal(result, expected)
# Case-insensitive
pat = values.dtype.type(".*bAd[_]+.*bad")
result = values.str.match(pat, case=False)
expected = xr.DataArray([True, False])
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.match(re.compile(pat, flags=re.IGNORECASE))
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_empty_str_methods() -> None:
empty = xr.DataArray(np.empty(shape=(0,), dtype="U"))
empty_str = empty
empty_int = xr.DataArray(np.empty(shape=(0,), dtype=int))
empty_bool = xr.DataArray(np.empty(shape=(0,), dtype=bool))
empty_bytes = xr.DataArray(np.empty(shape=(0,), dtype="S"))
# TODO: Determine why U and S dtype sizes don't match and figure
# out a reliable way to predict what they should be
assert empty_bool.dtype == empty.str.contains("a").dtype
assert empty_bool.dtype == empty.str.endswith("a").dtype
assert empty_bool.dtype == empty.str.match("^a").dtype
assert empty_bool.dtype == empty.str.startswith("a").dtype
assert empty_bool.dtype == empty.str.isalnum().dtype
assert empty_bool.dtype == empty.str.isalpha().dtype
assert empty_bool.dtype == empty.str.isdecimal().dtype
assert empty_bool.dtype == empty.str.isdigit().dtype
assert empty_bool.dtype == empty.str.islower().dtype
assert empty_bool.dtype == empty.str.isnumeric().dtype
assert empty_bool.dtype == empty.str.isspace().dtype
assert empty_bool.dtype == empty.str.istitle().dtype
assert empty_bool.dtype == empty.str.isupper().dtype
assert empty_bytes.dtype.kind == empty.str.encode("ascii").dtype.kind
assert empty_int.dtype.kind == empty.str.count("a").dtype.kind
assert empty_int.dtype.kind == empty.str.find("a").dtype.kind
assert empty_int.dtype.kind == empty.str.len().dtype.kind
assert empty_int.dtype.kind == empty.str.rfind("a").dtype.kind
assert empty_str.dtype.kind == empty.str.capitalize().dtype.kind
assert empty_str.dtype.kind == empty.str.center(42).dtype.kind
assert empty_str.dtype.kind == empty.str.get(0).dtype.kind
assert empty_str.dtype.kind == empty.str.lower().dtype.kind
assert empty_str.dtype.kind == empty.str.lstrip().dtype.kind
assert empty_str.dtype.kind == empty.str.pad(42).dtype.kind
assert empty_str.dtype.kind == empty.str.repeat(3).dtype.kind
assert empty_str.dtype.kind == empty.str.rstrip().dtype.kind
assert empty_str.dtype.kind == empty.str.slice(step=1).dtype.kind
assert empty_str.dtype.kind == empty.str.slice(stop=1).dtype.kind
assert empty_str.dtype.kind == empty.str.strip().dtype.kind
assert empty_str.dtype.kind == empty.str.swapcase().dtype.kind
assert empty_str.dtype.kind == empty.str.title().dtype.kind
assert empty_str.dtype.kind == empty.str.upper().dtype.kind
assert empty_str.dtype.kind == empty.str.wrap(42).dtype.kind
assert empty_str.dtype.kind == empty_bytes.str.decode("ascii").dtype.kind
assert_equal(empty_bool, empty.str.contains("a"))
assert_equal(empty_bool, empty.str.endswith("a"))
assert_equal(empty_bool, empty.str.match("^a"))
assert_equal(empty_bool, empty.str.startswith("a"))
assert_equal(empty_bool, empty.str.isalnum())
assert_equal(empty_bool, empty.str.isalpha())
assert_equal(empty_bool, empty.str.isdecimal())
assert_equal(empty_bool, empty.str.isdigit())
assert_equal(empty_bool, empty.str.islower())
assert_equal(empty_bool, empty.str.isnumeric())
assert_equal(empty_bool, empty.str.isspace())
assert_equal(empty_bool, empty.str.istitle())
assert_equal(empty_bool, empty.str.isupper())
assert_equal(empty_bytes, empty.str.encode("ascii"))
assert_equal(empty_int, empty.str.count("a"))
assert_equal(empty_int, empty.str.find("a"))
assert_equal(empty_int, empty.str.len())
assert_equal(empty_int, empty.str.rfind("a"))
assert_equal(empty_str, empty.str.capitalize())
assert_equal(empty_str, empty.str.center(42))
assert_equal(empty_str, empty.str.get(0))
assert_equal(empty_str, empty.str.lower())
assert_equal(empty_str, empty.str.lstrip())
assert_equal(empty_str, empty.str.pad(42))
assert_equal(empty_str, empty.str.repeat(3))
assert_equal(empty_str, empty.str.replace("a", "b"))
assert_equal(empty_str, empty.str.rstrip())
assert_equal(empty_str, empty.str.slice(step=1))
assert_equal(empty_str, empty.str.slice(stop=1))
assert_equal(empty_str, empty.str.strip())
assert_equal(empty_str, empty.str.swapcase())
assert_equal(empty_str, empty.str.title())
assert_equal(empty_str, empty.str.upper())
assert_equal(empty_str, empty.str.wrap(42))
assert_equal(empty_str, empty_bytes.str.decode("ascii"))
table = str.maketrans("a", "b")
assert empty_str.dtype.kind == empty.str.translate(table).dtype.kind
assert_equal(empty_str, empty.str.translate(table))
@pytest.mark.parametrize(
["func", "expected"],
[
pytest.param(
lambda x: x.str.isalnum(),
[True, True, True, True, True, False, True, True, False, False],
id="isalnum",
),
pytest.param(
lambda x: x.str.isalpha(),
[True, True, True, False, False, False, True, False, False, False],
id="isalpha",
),
pytest.param(
lambda x: x.str.isdigit(),
[False, False, False, True, False, False, False, True, False, False],
id="isdigit",
),
pytest.param(
lambda x: x.str.islower(),
[False, True, False, False, False, False, False, False, False, False],
id="islower",
),
pytest.param(
lambda x: x.str.isspace(),
[False, False, False, False, False, False, False, False, False, True],
id="isspace",
),
pytest.param(
lambda x: x.str.istitle(),
[True, False, True, False, True, False, False, False, False, False],
id="istitle",
),
pytest.param(
lambda x: x.str.isupper(),
[True, False, False, False, True, False, True, False, False, False],
id="isupper",
),
],
)
def test_ismethods(
dtype, func: Callable[[xr.DataArray], xr.DataArray], expected: list[bool]
) -> None:
values = xr.DataArray(
["A", "b", "Xy", "4", "3A", "", "TT", "55", "-", " "]
).astype(dtype)
expected_da = xr.DataArray(expected)
actual = func(values)
assert actual.dtype == expected_da.dtype
assert_equal(actual, expected_da)
def test_isnumeric() -> None:
# 0x00bc: ¼ VULGAR FRACTION ONE QUARTER
# 0x2605: ★ not number
# 0x1378: ፸ ETHIOPIC NUMBER SEVENTY
# 0xFF13: ３ Em 3
values = xr.DataArray(["A", "3", "¼", "★", "፸", "３", "four"])
exp_numeric = xr.DataArray([False, True, True, False, True, True, False])
exp_decimal = xr.DataArray([False, True, False, False, False, True, False])
res_numeric = values.str.isnumeric()
res_decimal = values.str.isdecimal()
assert res_numeric.dtype == exp_numeric.dtype
assert res_decimal.dtype == exp_decimal.dtype
assert_equal(res_numeric, exp_numeric)
assert_equal(res_decimal, exp_decimal)
def test_len(dtype) -> None:
values = ["foo", "fooo", "fooooo", "fooooooo"]
result = xr.DataArray(values).astype(dtype).str.len()
expected = xr.DataArray([len(x) for x in values])
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_find(dtype) -> None:
values = xr.DataArray(["ABCDEFG", "BCDEFEF", "DEFGHIJEF", "EFGHEF", "XXX"])
values = values.astype(dtype)
result_0 = values.str.find("EF")
result_1 = values.str.find("EF", side="left")
expected_0 = xr.DataArray([4, 3, 1, 0, -1])
expected_1 = xr.DataArray([v.find(dtype("EF")) for v in values.values])
assert result_0.dtype == expected_0.dtype
assert result_0.dtype == expected_1.dtype
assert result_1.dtype == expected_0.dtype
assert result_1.dtype == expected_1.dtype
assert_equal(result_0, expected_0)
assert_equal(result_0, expected_1)
assert_equal(result_1, expected_0)
assert_equal(result_1, expected_1)
result_0 = values.str.rfind("EF")
result_1 = values.str.find("EF", side="right")
expected_0 = xr.DataArray([4, 5, 7, 4, -1])
expected_1 = xr.DataArray([v.rfind(dtype("EF")) for v in values.values])
assert result_0.dtype == expected_0.dtype
assert result_0.dtype == expected_1.dtype
assert result_1.dtype == expected_0.dtype
assert result_1.dtype == expected_1.dtype
assert_equal(result_0, expected_0)
assert_equal(result_0, expected_1)
assert_equal(result_1, expected_0)
assert_equal(result_1, expected_1)
result_0 = values.str.find("EF", 3)
result_1 = values.str.find("EF", 3, side="left")
expected_0 = xr.DataArray([4, 3, 7, 4, -1])
expected_1 = xr.DataArray([v.find(dtype("EF"), 3) for v in values.values])
assert result_0.dtype == expected_0.dtype
assert result_0.dtype == expected_1.dtype
assert result_1.dtype == expected_0.dtype
assert result_1.dtype == expected_1.dtype
assert_equal(result_0, expected_0)
assert_equal(result_0, expected_1)
assert_equal(result_1, expected_0)
assert_equal(result_1, expected_1)
result_0 = values.str.rfind("EF", 3)
result_1 = values.str.find("EF", 3, side="right")
expected_0 = xr.DataArray([4, 5, 7, 4, -1])
expected_1 = xr.DataArray([v.rfind(dtype("EF"), 3) for v in values.values])
assert result_0.dtype == expected_0.dtype
assert result_0.dtype == expected_1.dtype
assert result_1.dtype == expected_0.dtype
assert result_1.dtype == expected_1.dtype
assert_equal(result_0, expected_0)
assert_equal(result_0, expected_1)
assert_equal(result_1, expected_0)
assert_equal(result_1, expected_1)
result_0 = values.str.find("EF", 3, 6)
result_1 = values.str.find("EF", 3, 6, side="left")
expected_0 = xr.DataArray([4, 3, -1, 4, -1])
expected_1 = xr.DataArray([v.find(dtype("EF"), 3, 6) for v in values.values])
assert result_0.dtype == expected_0.dtype
assert result_0.dtype == expected_1.dtype
assert result_1.dtype == expected_0.dtype
assert result_1.dtype == expected_1.dtype
assert_equal(result_0, expected_0)
assert_equal(result_0, expected_1)
assert_equal(result_1, expected_0)
assert_equal(result_1, expected_1)
result_0 = values.str.rfind("EF", 3, 6)
result_1 = values.str.find("EF", 3, 6, side="right")
expected_0 = xr.DataArray([4, 3, -1, 4, -1])
expected_1 = xr.DataArray([v.rfind(dtype("EF"), 3, 6) for v in values.values])
assert result_0.dtype == expected_0.dtype
assert result_0.dtype == expected_1.dtype
assert result_1.dtype == expected_0.dtype
assert result_1.dtype == expected_1.dtype
assert_equal(result_0, expected_0)
assert_equal(result_0, expected_1)
assert_equal(result_1, expected_0)
assert_equal(result_1, expected_1)
def test_find_broadcast(dtype) -> None:
values = xr.DataArray(
["ABCDEFG", "BCDEFEF", "DEFGHIJEF", "EFGHEF", "XXX"], dims=["X"]
)
values = values.astype(dtype)
sub = xr.DataArray(["EF", "BC", "XX"], dims=["Y"]).astype(dtype)
start = xr.DataArray([0, 7], dims=["Z"])
end = xr.DataArray([6, 9], dims=["Z"])
result_0 = values.str.find(sub, start, end)
result_1 = values.str.find(sub, start, end, side="left")
expected = xr.DataArray(
[
[[4, -1], [1, -1], [-1, -1]],
[[3, -1], [0, -1], [-1, -1]],
[[1, 7], [-1, -1], [-1, -1]],
[[0, -1], [-1, -1], [-1, -1]],
[[-1, -1], [-1, -1], [0, -1]],
],
dims=["X", "Y", "Z"],
)
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = values.str.rfind(sub, start, end)
result_1 = values.str.find(sub, start, end, side="right")
expected = xr.DataArray(
[
[[4, -1], [1, -1], [-1, -1]],
[[3, -1], [0, -1], [-1, -1]],
[[1, 7], [-1, -1], [-1, -1]],
[[4, -1], [-1, -1], [-1, -1]],
[[-1, -1], [-1, -1], [1, -1]],
],
dims=["X", "Y", "Z"],
)
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
def test_index(dtype) -> None:
s = xr.DataArray(["ABCDEFG", "BCDEFEF", "DEFGHIJEF", "EFGHEF"]).astype(dtype)
result_0 = s.str.index("EF")
result_1 = s.str.index("EF", side="left")
expected = xr.DataArray([4, 3, 1, 0])
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = s.str.rindex("EF")
result_1 = s.str.index("EF", side="right")
expected = xr.DataArray([4, 5, 7, 4])
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = s.str.index("EF", 3)
result_1 = s.str.index("EF", 3, side="left")
expected = xr.DataArray([4, 3, 7, 4])
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = s.str.rindex("EF", 3)
result_1 = s.str.index("EF", 3, side="right")
expected = xr.DataArray([4, 5, 7, 4])
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = s.str.index("E", 4, 8)
result_1 = s.str.index("E", 4, 8, side="left")
expected = xr.DataArray([4, 5, 7, 4])
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = s.str.rindex("E", 0, 5)
result_1 = s.str.index("E", 0, 5, side="right")
expected = xr.DataArray([4, 3, 1, 4])
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
matchtype = "subsection" if dtype == np.bytes_ else "substring"
with pytest.raises(ValueError, match=f"{matchtype} not found"):
s.str.index("DE")
def test_index_broadcast(dtype) -> None:
values = xr.DataArray(
["ABCDEFGEFDBCA", "BCDEFEFEFDBC", "DEFBCGHIEFBC", "EFGHBCEFBCBCBCEF"],
dims=["X"],
)
values = values.astype(dtype)
sub = xr.DataArray(["EF", "BC"], dims=["Y"]).astype(dtype)
start = xr.DataArray([0, 6], dims=["Z"])
end = xr.DataArray([6, 12], dims=["Z"])
result_0 = values.str.index(sub, start, end)
result_1 = values.str.index(sub, start, end, side="left")
expected = xr.DataArray(
[[[4, 7], [1, 10]], [[3, 7], [0, 10]], [[1, 8], [3, 10]], [[0, 6], [4, 8]]],
dims=["X", "Y", "Z"],
)
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
result_0 = values.str.rindex(sub, start, end)
result_1 = values.str.index(sub, start, end, side="right")
expected = xr.DataArray(
[[[4, 7], [1, 10]], [[3, 7], [0, 10]], [[1, 8], [3, 10]], [[0, 6], [4, 10]]],
dims=["X", "Y", "Z"],
)
assert result_0.dtype == expected.dtype
assert result_1.dtype == expected.dtype
assert_equal(result_0, expected)
assert_equal(result_1, expected)
def test_translate() -> None:
values = xr.DataArray(["abcdefg", "abcc", "cdddfg", "cdefggg"])
table = str.maketrans("abc", "cde")
result = values.str.translate(table)
expected = xr.DataArray(["cdedefg", "cdee", "edddfg", "edefggg"])
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_pad_center_ljust_rjust(dtype) -> None:
values = xr.DataArray(["a", "b", "c", "eeeee"]).astype(dtype)
result = values.str.center(5)
expected = xr.DataArray([" a ", " b ", " c ", "eeeee"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.pad(5, side="both")
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.ljust(5)
expected = xr.DataArray(["a ", "b ", "c ", "eeeee"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.pad(5, side="right")
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.rjust(5)
expected = xr.DataArray([" a", " b", " c", "eeeee"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.pad(5, side="left")
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_pad_center_ljust_rjust_fillchar(dtype) -> None:
values = xr.DataArray(["a", "bb", "cccc", "ddddd", "eeeeee"]).astype(dtype)
result = values.str.center(5, fillchar="X")
expected = xr.DataArray(["XXaXX", "XXbbX", "Xcccc", "ddddd", "eeeeee"]).astype(
dtype
)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.pad(5, side="both", fillchar="X")
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.ljust(5, fillchar="X")
expected = xr.DataArray(["aXXXX", "bbXXX", "ccccX", "ddddd", "eeeeee"]).astype(
dtype
)
assert result.dtype == expected.dtype
assert_equal(result, expected.astype(dtype))
result = values.str.pad(5, side="right", fillchar="X")
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.rjust(5, fillchar="X")
expected = xr.DataArray(["XXXXa", "XXXbb", "Xcccc", "ddddd", "eeeeee"]).astype(
dtype
)
assert result.dtype == expected.dtype
assert_equal(result, expected.astype(dtype))
result = values.str.pad(5, side="left", fillchar="X")
assert result.dtype == expected.dtype
assert_equal(result, expected)
# If fillchar is not a charatter, normal str raises TypeError
# 'aaa'.ljust(5, 'XY')
# TypeError: must be char, not str
template = "fillchar must be a character, not {dtype}"
with pytest.raises(TypeError, match=template.format(dtype="str")):
values.str.center(5, fillchar="XY")
with pytest.raises(TypeError, match=template.format(dtype="str")):
values.str.ljust(5, fillchar="XY")
with pytest.raises(TypeError, match=template.format(dtype="str")):
values.str.rjust(5, fillchar="XY")
with pytest.raises(TypeError, match=template.format(dtype="str")):
values.str.pad(5, fillchar="XY")
def test_pad_center_ljust_rjust_broadcast(dtype) -> None:
values = xr.DataArray(["a", "bb", "cccc", "ddddd", "eeeeee"], dims="X").astype(
dtype
)
width = xr.DataArray([5, 4], dims="Y")
fillchar = xr.DataArray(["X", "#"], dims="Y").astype(dtype)
result = values.str.center(width, fillchar=fillchar)
expected = xr.DataArray(
[
["XXaXX", "#a##"],
["XXbbX", "#bb#"],
["Xcccc", "cccc"],
["ddddd", "ddddd"],
["eeeeee", "eeeeee"],
],
dims=["X", "Y"],
).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.pad(width, side="both", fillchar=fillchar)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.ljust(width, fillchar=fillchar)
expected = xr.DataArray(
[
["aXXXX", "a###"],
["bbXXX", "bb##"],
["ccccX", "cccc"],
["ddddd", "ddddd"],
["eeeeee", "eeeeee"],
],
dims=["X", "Y"],
).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected.astype(dtype))
result = values.str.pad(width, side="right", fillchar=fillchar)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.rjust(width, fillchar=fillchar)
expected = xr.DataArray(
[
["XXXXa", "###a"],
["XXXbb", "##bb"],
["Xcccc", "cccc"],
["ddddd", "ddddd"],
["eeeeee", "eeeeee"],
],
dims=["X", "Y"],
).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected.astype(dtype))
result = values.str.pad(width, side="left", fillchar=fillchar)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_zfill(dtype) -> None:
values = xr.DataArray(["1", "22", "aaa", "333", "45678"]).astype(dtype)
result = values.str.zfill(5)
expected = xr.DataArray(["00001", "00022", "00aaa", "00333", "45678"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.zfill(3)
expected = xr.DataArray(["001", "022", "aaa", "333", "45678"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_zfill_broadcast(dtype) -> None:
values = xr.DataArray(["1", "22", "aaa", "333", "45678"]).astype(dtype)
width = np.array([4, 5, 0, 3, 8])
result = values.str.zfill(width)
expected = xr.DataArray(["0001", "00022", "aaa", "333", "00045678"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_slice(dtype) -> None:
arr = xr.DataArray(["aafootwo", "aabartwo", "aabazqux"]).astype(dtype)
result = arr.str.slice(2, 5)
exp = xr.DataArray(["foo", "bar", "baz"]).astype(dtype)
assert result.dtype == exp.dtype
assert_equal(result, exp)
for start, stop, step in [(0, 3, -1), (None, None, -1), (3, 10, 2), (3, 0, -1)]:
try:
result = arr.str[start:stop:step]
expected = xr.DataArray([s[start:stop:step] for s in arr.values])
assert_equal(result, expected.astype(dtype))
except IndexError:
print(f"failed on {start}:{stop}:{step}")
raise
def test_slice_broadcast(dtype) -> None:
arr = xr.DataArray(["aafootwo", "aabartwo", "aabazqux"]).astype(dtype)
start = xr.DataArray([1, 2, 3])
stop = 5
result = arr.str.slice(start=start, stop=stop)
exp = xr.DataArray(["afoo", "bar", "az"]).astype(dtype)
assert result.dtype == exp.dtype
assert_equal(result, exp)
def test_slice_replace(dtype) -> None:
da = lambda x: xr.DataArray(x).astype(dtype)
values = da(["short", "a bit longer", "evenlongerthanthat", ""])
expected = da(["shrt", "a it longer", "evnlongerthanthat", ""])
result = values.str.slice_replace(2, 3)
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["shzrt", "a zit longer", "evznlongerthanthat", "z"])
result = values.str.slice_replace(2, 3, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["shzort", "a zbit longer", "evzenlongerthanthat", "z"])
result = values.str.slice_replace(2, 2, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["shzort", "a zbit longer", "evzenlongerthanthat", "z"])
result = values.str.slice_replace(2, 1, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["shorz", "a bit longez", "evenlongerthanthaz", "z"])
result = values.str.slice_replace(-1, None, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["zrt", "zer", "zat", "z"])
result = values.str.slice_replace(None, -2, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["shortz", "a bit znger", "evenlozerthanthat", "z"])
result = values.str.slice_replace(6, 8, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
expected = da(["zrt", "a zit longer", "evenlongzerthanthat", "z"])
result = values.str.slice_replace(-10, 3, "z")
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_slice_replace_broadcast(dtype) -> None:
values = xr.DataArray(["short", "a bit longer", "evenlongerthanthat", ""]).astype(
dtype
)
start = 2
stop = np.array([4, 5, None, 7])
repl = "test"
expected = xr.DataArray(["shtestt", "a test longer", "evtest", "test"]).astype(
dtype
)
result = values.str.slice_replace(start, stop, repl)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_strip_lstrip_rstrip(dtype) -> None:
values = xr.DataArray([" aa ", " bb \n", "cc "]).astype(dtype)
result = values.str.strip()
expected = xr.DataArray(["aa", "bb", "cc"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.lstrip()
expected = xr.DataArray(["aa ", "bb \n", "cc "]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.rstrip()
expected = xr.DataArray([" aa", " bb", "cc"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_strip_lstrip_rstrip_args(dtype) -> None:
values = xr.DataArray(["xxABCxx", "xx BNSD", "LDFJH xx"]).astype(dtype)
result = values.str.strip("x")
expected = xr.DataArray(["ABC", " BNSD", "LDFJH "]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.lstrip("x")
expected = xr.DataArray(["ABCxx", " BNSD", "LDFJH xx"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.rstrip("x")
expected = xr.DataArray(["xxABC", "xx BNSD", "LDFJH "]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_strip_lstrip_rstrip_broadcast(dtype) -> None:
values = xr.DataArray(["xxABCxx", "yy BNSD", "LDFJH zz"]).astype(dtype)
to_strip = xr.DataArray(["x", "y", "z"]).astype(dtype)
result = values.str.strip(to_strip)
expected = xr.DataArray(["ABC", " BNSD", "LDFJH "]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.lstrip(to_strip)
expected = xr.DataArray(["ABCxx", " BNSD", "LDFJH zz"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.rstrip(to_strip)
expected = xr.DataArray(["xxABC", "yy BNSD", "LDFJH "]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_wrap() -> None:
# test values are: two words less than width, two words equal to width,
# two words greater than width, one word less than width, one word
# equal to width, one word greater than width, multiple tokens with
# trailing whitespace equal to width
values = xr.DataArray(
[
"hello world",
"hello world!",
"hello world!!",
"abcdefabcde",
"abcdefabcdef",
"abcdefabcdefa",
"ab ab ab ab ",
"ab ab ab ab a",
"\t",
]
)
# expected values
expected = xr.DataArray(
[
"hello world",
"hello world!",
"hello\nworld!!",
"abcdefabcde",
"abcdefabcdef",
"abcdefabcdef\na",
"ab ab ab ab",
"ab ab ab ab\na",
"",
]
)
result = values.str.wrap(12, break_long_words=True)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# test with pre and post whitespace (non-unicode), NaN, and non-ascii
# Unicode
values = xr.DataArray([" pre ", "\xac\u20ac\U00008000 abadcafe"])
expected = xr.DataArray([" pre", "\xac\u20ac\U00008000 ab\nadcafe"])
result = values.str.wrap(6)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_wrap_kwargs_passed() -> None:
# GH4334
values = xr.DataArray(" hello world ")
result = values.str.wrap(7)
expected = xr.DataArray(" hello\nworld")
assert result.dtype == expected.dtype
assert_equal(result, expected)
result = values.str.wrap(7, drop_whitespace=False)
expected = xr.DataArray(" hello\n world\n ")
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_get(dtype) -> None:
values = xr.DataArray(["a_b_c", "c_d_e", "f_g_h"]).astype(dtype)
result = values.str[2]
expected = xr.DataArray(["b", "d", "g"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# bounds testing
values = xr.DataArray(["1_2_3_4_5", "6_7_8_9_10", "11_12"]).astype(dtype)
# positive index
result = values.str[5]
expected = xr.DataArray(["_", "_", ""]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
# negative index
result = values.str[-6]
expected = xr.DataArray(["_", "8", ""]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_get_default(dtype) -> None:
# GH4334
values = xr.DataArray(["a_b", "c", ""]).astype(dtype)
result = values.str.get(2, "default")
expected = xr.DataArray(["b", "default", "default"]).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_get_broadcast(dtype) -> None:
values = xr.DataArray(["a_b_c", "c_d_e", "f_g_h"], dims=["X"]).astype(dtype)
inds = xr.DataArray([0, 2], dims=["Y"])
result = values.str.get(inds)
expected = xr.DataArray(
[["a", "b"], ["c", "d"], ["f", "g"]], dims=["X", "Y"]
).astype(dtype)
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_encode_decode() -> None:
data = xr.DataArray(["a", "b", "a\xe4"])
encoded = data.str.encode("utf-8")
decoded = encoded.str.decode("utf-8")
assert data.dtype == decoded.dtype
assert_equal(data, decoded)
def test_encode_decode_errors() -> None:
encodeBase = xr.DataArray(["a", "b", "a\x9d"])
msg = (
r"'charmap' codec can't encode character '\\x9d' in position 1:"
" character maps to <undefined>"
)
with pytest.raises(UnicodeEncodeError, match=msg):
encodeBase.str.encode("cp1252")
f = lambda x: x.encode("cp1252", "ignore")
result = encodeBase.str.encode("cp1252", "ignore")
expected = xr.DataArray([f(x) for x in encodeBase.values.tolist()])
assert result.dtype == expected.dtype
assert_equal(result, expected)
decodeBase = xr.DataArray([b"a", b"b", b"a\x9d"])
msg = (
"'charmap' codec can't decode byte 0x9d in position 1:"
" character maps to <undefined>"
)
with pytest.raises(UnicodeDecodeError, match=msg):
decodeBase.str.decode("cp1252")
f = lambda x: x.decode("cp1252", "ignore")
result = decodeBase.str.decode("cp1252", "ignore")
expected = xr.DataArray([f(x) for x in decodeBase.values.tolist()])
assert result.dtype == expected.dtype
assert_equal(result, expected)
def test_partition_whitespace(dtype) -> None:
values = xr.DataArray(
[
["abc def", "spam eggs swallow", "red_blue"],
["test0 test1 test2 test3", "", "abra ka da bra"],
],
dims=["X", "Y"],
).astype(dtype)
exp_part_dim_list = [
[
["abc", " ", "def"],
["spam", " ", "eggs swallow"],
["red_blue", "", ""],
],
[
["test0", " ", "test1 test2 test3"],
["", "", ""],
["abra", " ", "ka da bra"],
],
]
exp_rpart_dim_list = [
[
["abc", " ", "def"],
["spam eggs", " ", "swallow"],
["", "", "red_blue"],
],
[
["test0 test1 test2", " ", "test3"],
["", "", ""],
["abra ka da", " ", "bra"],
],
]
exp_part_dim = xr.DataArray(exp_part_dim_list, dims=["X", "Y", "ZZ"]).astype(dtype)
exp_rpart_dim = xr.DataArray(exp_rpart_dim_list, dims=["X", "Y", "ZZ"]).astype(
dtype
)
res_part_dim = values.str.partition(dim="ZZ")
res_rpart_dim = values.str.rpartition(dim="ZZ")
assert res_part_dim.dtype == exp_part_dim.dtype
assert res_rpart_dim.dtype == exp_rpart_dim.dtype
assert_equal(res_part_dim, exp_part_dim)
assert_equal(res_rpart_dim, exp_rpart_dim)
def test_partition_comma(dtype) -> None:
values = xr.DataArray(
[
["abc, def", "spam, eggs, swallow", "red_blue"],
["test0, test1, test2, test3", "", "abra, ka, da, bra"],
],
dims=["X", "Y"],
).astype(dtype)
exp_part_dim_list = [
[
["abc", ", ", "def"],
["spam", ", ", "eggs, swallow"],
["red_blue", "", ""],
],
[
["test0", ", ", "test1, test2, test3"],
["", "", ""],
["abra", ", ", "ka, da, bra"],
],
]
exp_rpart_dim_list = [
[
["abc", ", ", "def"],
["spam, eggs", ", ", "swallow"],
["", "", "red_blue"],
],
[
["test0, test1, test2", ", ", "test3"],
["", "", ""],
["abra, ka, da", ", ", "bra"],
],
]
exp_part_dim = xr.DataArray(exp_part_dim_list, dims=["X", "Y", "ZZ"]).astype(dtype)
exp_rpart_dim = xr.DataArray(exp_rpart_dim_list, dims=["X", "Y", "ZZ"]).astype(
dtype
)
res_part_dim = values.str.partition(sep=", ", dim="ZZ")
res_rpart_dim = values.str.rpartition(sep=", ", dim="ZZ")
assert res_part_dim.dtype == exp_part_dim.dtype
assert res_rpart_dim.dtype == exp_rpart_dim.dtype
assert_equal(res_part_dim, exp_part_dim)
assert_equal(res_rpart_dim, exp_rpart_dim)
def test_partition_empty(dtype) -> None:
values = xr.DataArray([], dims=["X"]).astype(dtype)
expected = xr.DataArray(np.zeros((0, 0)), dims=["X", "ZZ"]).astype(dtype)
res = values.str.partition(sep=", ", dim="ZZ")
assert res.dtype == expected.dtype
assert_equal(res, expected)
@pytest.mark.parametrize(
["func", "expected"],
[
pytest.param(
lambda x: x.str.split(dim=None),
[
[["abc", "def"], ["spam", "eggs", "swallow"], ["red_blue"]],
[["test0", "test1", "test2", "test3"], [], ["abra", "ka", "da", "bra"]],
],
id="split_full",
),
pytest.param(
lambda x: x.str.rsplit(dim=None),
[
[["abc", "def"], ["spam", "eggs", "swallow"], ["red_blue"]],
[["test0", "test1", "test2", "test3"], [], ["abra", "ka", "da", "bra"]],
],
id="rsplit_full",
),
pytest.param(
lambda x: x.str.split(dim=None, maxsplit=1),
[
[["abc", "def"], ["spam", "eggs\tswallow"], ["red_blue"]],
[["test0", "test1\ntest2\n\ntest3"], [], ["abra", "ka\nda\tbra"]],
],
id="split_1",
),
pytest.param(
lambda x: x.str.rsplit(dim=None, maxsplit=1),
[
[["abc", "def"], ["spam\t\teggs", "swallow"], ["red_blue"]],
[["test0\ntest1\ntest2", "test3"], [], ["abra ka\nda", "bra"]],
],
id="rsplit_1",
),
],
)
def test_split_whitespace_nodim(
dtype, func: Callable[[xr.DataArray], xr.DataArray], expected: xr.DataArray
) -> None:
values = xr.DataArray(
[
["abc def", "spam\t\teggs\tswallow", "red_blue"],
["test0\ntest1\ntest2\n\ntest3", "", "abra ka\nda\tbra"],
],
dims=["X", "Y"],
).astype(dtype)
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected_da = xr.DataArray(expected_np, dims=["X", "Y"])
actual = func(values)
assert actual.dtype == expected_da.dtype
assert_equal(actual, expected_da)
@pytest.mark.parametrize(
["func", "expected"],
[
pytest.param(
lambda x: x.str.split(dim="ZZ"),
[
[
["abc", "def", "", ""],
["spam", "eggs", "swallow", ""],
["red_blue", "", "", ""],
],
[
["test0", "test1", "test2", "test3"],
["", "", "", ""],
["abra", "ka", "da", "bra"],
],
],
id="split_full",
),
pytest.param(
lambda x: x.str.rsplit(dim="ZZ"),
[
[
["", "", "abc", "def"],
["", "spam", "eggs", "swallow"],
["", "", "", "red_blue"],
],
[
["test0", "test1", "test2", "test3"],
["", "", "", ""],
["abra", "ka", "da", "bra"],
],
],
id="rsplit_full",
),
pytest.param(
lambda x: x.str.split(dim="ZZ", maxsplit=1),
[
[["abc", "def"], ["spam", "eggs\tswallow"], ["red_blue", ""]],
[["test0", "test1\ntest2\n\ntest3"], ["", ""], ["abra", "ka\nda\tbra"]],
],
id="split_1",
),
pytest.param(
lambda x: x.str.rsplit(dim="ZZ", maxsplit=1),
[
[["abc", "def"], ["spam\t\teggs", "swallow"], ["", "red_blue"]],
[["test0\ntest1\ntest2", "test3"], ["", ""], ["abra ka\nda", "bra"]],
],
id="rsplit_1",
),
],
)
def test_split_whitespace_dim(
dtype, func: Callable[[xr.DataArray], xr.DataArray], expected: xr.DataArray
) -> None:
values = xr.DataArray(
[
["abc def", "spam\t\teggs\tswallow", "red_blue"],
["test0\ntest1\ntest2\n\ntest3", "", "abra ka\nda\tbra"],
],
dims=["X", "Y"],
).astype(dtype)
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected_da = xr.DataArray(expected_np, dims=["X", "Y", "ZZ"]).astype(dtype)
actual = func(values)
assert actual.dtype == expected_da.dtype
assert_equal(actual, expected_da)
@pytest.mark.parametrize(
["func", "expected"],
[
pytest.param(
lambda x: x.str.split(sep=",", dim=None),
[
[["abc", "def"], ["spam", "", "eggs", "swallow"], ["red_blue"]],
[
["test0", "test1", "test2", "test3"],
[""],
["abra", "ka", "da", "bra"],
],
],
id="split_full",
),
pytest.param(
lambda x: x.str.rsplit(sep=",", dim=None),
[
[["abc", "def"], ["spam", "", "eggs", "swallow"], ["red_blue"]],
[
["test0", "test1", "test2", "test3"],
[""],
["abra", "ka", "da", "bra"],
],
],
id="rsplit_full",
),
pytest.param(
lambda x: x.str.split(sep=",", dim=None, maxsplit=1),
[
[["abc", "def"], ["spam", ",eggs,swallow"], ["red_blue"]],
[["test0", "test1,test2,test3"], [""], ["abra", "ka,da,bra"]],
],
id="split_1",
),
pytest.param(
lambda x: x.str.rsplit(sep=",", dim=None, maxsplit=1),
[
[["abc", "def"], ["spam,,eggs", "swallow"], ["red_blue"]],
[["test0,test1,test2", "test3"], [""], ["abra,ka,da", "bra"]],
],
id="rsplit_1",
),
pytest.param(
lambda x: x.str.split(sep=",", dim=None, maxsplit=10),
[
[["abc", "def"], ["spam", "", "eggs", "swallow"], ["red_blue"]],
[
["test0", "test1", "test2", "test3"],
[""],
["abra", "ka", "da", "bra"],
],
],
id="split_10",
),
pytest.param(
lambda x: x.str.rsplit(sep=",", dim=None, maxsplit=10),
[
[["abc", "def"], ["spam", "", "eggs", "swallow"], ["red_blue"]],
[
["test0", "test1", "test2", "test3"],
[""],
["abra", "ka", "da", "bra"],
],
],
id="rsplit_10",
),
],
)
def test_split_comma_nodim(
dtype, func: Callable[[xr.DataArray], xr.DataArray], expected: xr.DataArray
) -> None:
values = xr.DataArray(
[
["abc,def", "spam,,eggs,swallow", "red_blue"],
["test0,test1,test2,test3", "", "abra,ka,da,bra"],
],
dims=["X", "Y"],
).astype(dtype)
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected_da = xr.DataArray(expected_np, dims=["X", "Y"])
actual = func(values)
assert actual.dtype == expected_da.dtype
assert_equal(actual, expected_da)
@pytest.mark.parametrize(
["func", "expected"],
[
pytest.param(
lambda x: x.str.split(sep=",", dim="ZZ"),
[
[
["abc", "def", "", ""],
["spam", "", "eggs", "swallow"],
["red_blue", "", "", ""],
],
[
["test0", "test1", "test2", "test3"],
["", "", "", ""],
["abra", "ka", "da", "bra"],
],
],
id="split_full",
),
pytest.param(
lambda x: x.str.rsplit(sep=",", dim="ZZ"),
[
[
["", "", "abc", "def"],
["spam", "", "eggs", "swallow"],
["", "", "", "red_blue"],
],
[
["test0", "test1", "test2", "test3"],
["", "", "", ""],
["abra", "ka", "da", "bra"],
],
],
id="rsplit_full",
),
pytest.param(
lambda x: x.str.split(sep=",", dim="ZZ", maxsplit=1),
[
[["abc", "def"], ["spam", ",eggs,swallow"], ["red_blue", ""]],
[["test0", "test1,test2,test3"], ["", ""], ["abra", "ka,da,bra"]],
],
id="split_1",
),
pytest.param(
lambda x: x.str.rsplit(sep=",", dim="ZZ", maxsplit=1),
[
[["abc", "def"], ["spam,,eggs", "swallow"], ["", "red_blue"]],
[["test0,test1,test2", "test3"], ["", ""], ["abra,ka,da", "bra"]],
],
id="rsplit_1",
),
pytest.param(
lambda x: x.str.split(sep=",", dim="ZZ", maxsplit=10),
[
[
["abc", "def", "", ""],
["spam", "", "eggs", "swallow"],
["red_blue", "", "", ""],
],
[
["test0", "test1", "test2", "test3"],
["", "", "", ""],
["abra", "ka", "da", "bra"],
],
],
id="split_10",
),
pytest.param(
lambda x: x.str.rsplit(sep=",", dim="ZZ", maxsplit=10),
[
[
["", "", "abc", "def"],
["spam", "", "eggs", "swallow"],
["", "", "", "red_blue"],
],
[
["test0", "test1", "test2", "test3"],
["", "", "", ""],
["abra", "ka", "da", "bra"],
],
],
id="rsplit_10",
),
],
)
def test_split_comma_dim(
dtype, func: Callable[[xr.DataArray], xr.DataArray], expected: xr.DataArray
) -> None:
values = xr.DataArray(
[
["abc,def", "spam,,eggs,swallow", "red_blue"],
["test0,test1,test2,test3", "", "abra,ka,da,bra"],
],
dims=["X", "Y"],
).astype(dtype)
expected_dtype = [[[dtype(x) for x in y] for y in z] for z in expected]
expected_np = np.array(expected_dtype, dtype=np.object_)
expected_da = xr.DataArray(expected_np, dims=["X", "Y", "ZZ"]).astype(dtype)
actual = func(values)
assert actual.dtype == expected_da.dtype
assert_equal(actual, expected_da)
def test_splitters_broadcast(dtype) -> None:
values = xr.DataArray(
["ab cd,de fg", "spam, ,eggs swallow", "red_blue"],
dims=["X"],
).astype(dtype)
sep = xr.DataArray(
[" ", ","],
dims=["Y"],
).astype(dtype)
expected_left = xr.DataArray(
[
[["ab", "cd,de fg"], ["ab cd", "de fg"]],
[["spam,", ",eggs swallow"], ["spam", " ,eggs swallow"]],
[["red_blue", ""], ["red_blue", ""]],
],
dims=["X", "Y", "ZZ"],
).astype(dtype)
expected_right = xr.DataArray(
[
[["ab cd,de", "fg"], ["ab cd", "de fg"]],
[["spam, ,eggs", "swallow"], ["spam, ", "eggs swallow"]],
[["", "red_blue"], ["", "red_blue"]],
],
dims=["X", "Y", "ZZ"],
).astype(dtype)
res_left = values.str.split(dim="ZZ", sep=sep, maxsplit=1)
res_right = values.str.rsplit(dim="ZZ", sep=sep, maxsplit=1)
# assert res_left.dtype == expected_left.dtype
# assert res_right.dtype == expected_right.dtype
assert_equal(res_left, expected_left)
assert_equal(res_right, expected_right)
expected_left = xr.DataArray(
[
[["ab", " ", "cd,de fg"], ["ab cd", ",", "de fg"]],
[["spam,", " ", ",eggs swallow"], ["spam", ",", " ,eggs swallow"]],
[["red_blue", "", ""], ["red_blue", "", ""]],
],
dims=["X", "Y", "ZZ"],
).astype(dtype)
expected_right = xr.DataArray(
[
[["ab", " ", "cd,de fg"], ["ab cd", ",", "de fg"]],
[["spam,", " ", ",eggs swallow"], ["spam", ",", " ,eggs swallow"]],
[["red_blue", "", ""], ["red_blue", "", ""]],
],
dims=["X", "Y", "ZZ"],
).astype(dtype)
res_left = values.str.partition(dim="ZZ", sep=sep)
res_right = values.str.partition(dim="ZZ", sep=sep)
# assert res_left.dtype == expected_left.dtype
# assert res_right.dtype == expected_right.dtype
assert_equal(res_left, expected_left)
assert_equal(res_right, expected_right)
def test_split_empty(dtype) -> None:
values = xr.DataArray([], dims=["X"]).astype(dtype)
expected = xr.DataArray(np.zeros((0, 0)), dims=["X", "ZZ"]).astype(dtype)
res = values.str.split(sep=", ", dim="ZZ")
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_get_dummies(dtype) -> None:
values_line = xr.DataArray(
[["a|ab~abc|abc", "ab", "a||abc|abcd"], ["abcd|ab|a", "abc|ab~abc", "|a"]],
dims=["X", "Y"],
).astype(dtype)
values_comma = xr.DataArray(
[["a~ab|abc~~abc", "ab", "a~abc~abcd"], ["abcd~ab~a", "abc~ab|abc", "~a"]],
dims=["X", "Y"],
).astype(dtype)
vals_line = np.array(["a", "ab", "abc", "abcd", "ab~abc"]).astype(dtype)
vals_comma = np.array(["a", "ab", "abc", "abcd", "ab|abc"]).astype(dtype)
expected_list = [
[
[True, False, True, False, True],
[False, True, False, False, False],
[True, False, True, True, False],
],
[
[True, True, False, True, False],
[False, False, True, False, True],
[True, False, False, False, False],
],
]
expected_np = np.array(expected_list)
expected = xr.DataArray(expected_np, dims=["X", "Y", "ZZ"])
targ_line = expected.copy()
targ_comma = expected.copy()
targ_line.coords["ZZ"] = vals_line
targ_comma.coords["ZZ"] = vals_comma
res_default = values_line.str.get_dummies(dim="ZZ")
res_line = values_line.str.get_dummies(dim="ZZ", sep="|")
res_comma = values_comma.str.get_dummies(dim="ZZ", sep="~")
assert res_default.dtype == targ_line.dtype
assert res_line.dtype == targ_line.dtype
assert res_comma.dtype == targ_comma.dtype
assert_equal(res_default, targ_line)
assert_equal(res_line, targ_line)
assert_equal(res_comma, targ_comma)
def test_get_dummies_broadcast(dtype) -> None:
values = xr.DataArray(
["x~x|x~x", "x", "x|x~x", "x~x"],
dims=["X"],
).astype(dtype)
sep = xr.DataArray(
["|", "~"],
dims=["Y"],
).astype(dtype)
expected_list = [
[[False, False, True], [True, True, False]],
[[True, False, False], [True, False, False]],
[[True, False, True], [True, True, False]],
[[False, False, True], [True, False, False]],
]
expected_np = np.array(expected_list)
expected = xr.DataArray(expected_np, dims=["X", "Y", "ZZ"])
expected.coords["ZZ"] = np.array(["x", "x|x", "x~x"]).astype(dtype)
res = values.str.get_dummies(dim="ZZ", sep=sep)
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_get_dummies_empty(dtype) -> None:
values = xr.DataArray([], dims=["X"]).astype(dtype)
expected = xr.DataArray(np.zeros((0, 0)), dims=["X", "ZZ"]).astype(dtype)
res = values.str.get_dummies(dim="ZZ")
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_splitters_empty_str(dtype) -> None:
values = xr.DataArray(
[["", "", ""], ["", "", ""]],
dims=["X", "Y"],
).astype(dtype)
targ_partition_dim = xr.DataArray(
[
[["", "", ""], ["", "", ""], ["", "", ""]],
[["", "", ""], ["", "", ""], ["", "", ""]],
],
dims=["X", "Y", "ZZ"],
).astype(dtype)
targ_partition_none_list = [
[["", "", ""], ["", "", ""], ["", "", ""]],
[["", "", ""], ["", "", ""], ["", "", "", ""]],
]
targ_partition_none_list = [
[[dtype(x) for x in y] for y in z] for z in targ_partition_none_list
]
targ_partition_none_np = np.array(targ_partition_none_list, dtype=np.object_)
del targ_partition_none_np[-1, -1][-1]
targ_partition_none = xr.DataArray(
targ_partition_none_np,
dims=["X", "Y"],
)
targ_split_dim = xr.DataArray(
[[[""], [""], [""]], [[""], [""], [""]]],
dims=["X", "Y", "ZZ"],
).astype(dtype)
targ_split_none = xr.DataArray(
np.array([[[], [], []], [[], [], [""]]], dtype=np.object_),
dims=["X", "Y"],
)
del targ_split_none.data[-1, -1][-1]
res_partition_dim = values.str.partition(dim="ZZ")
res_rpartition_dim = values.str.rpartition(dim="ZZ")
res_partition_none = values.str.partition(dim=None)
res_rpartition_none = values.str.rpartition(dim=None)
res_split_dim = values.str.split(dim="ZZ")
res_rsplit_dim = values.str.rsplit(dim="ZZ")
res_split_none = values.str.split(dim=None)
res_rsplit_none = values.str.rsplit(dim=None)
res_dummies = values.str.rsplit(dim="ZZ")
assert res_partition_dim.dtype == targ_partition_dim.dtype
assert res_rpartition_dim.dtype == targ_partition_dim.dtype
assert res_partition_none.dtype == targ_partition_none.dtype
assert res_rpartition_none.dtype == targ_partition_none.dtype
assert res_split_dim.dtype == targ_split_dim.dtype
assert res_rsplit_dim.dtype == targ_split_dim.dtype
assert res_split_none.dtype == targ_split_none.dtype
assert res_rsplit_none.dtype == targ_split_none.dtype
assert res_dummies.dtype == targ_split_dim.dtype
assert_equal(res_partition_dim, targ_partition_dim)
assert_equal(res_rpartition_dim, targ_partition_dim)
assert_equal(res_partition_none, targ_partition_none)
assert_equal(res_rpartition_none, targ_partition_none)
assert_equal(res_split_dim, targ_split_dim)
assert_equal(res_rsplit_dim, targ_split_dim)
assert_equal(res_split_none, targ_split_none)
assert_equal(res_rsplit_none, targ_split_none)
assert_equal(res_dummies, targ_split_dim)
def test_cat_str(dtype) -> None:
values_1 = xr.DataArray(
[["a", "bb", "cccc"], ["ddddd", "eeee", "fff"]],
dims=["X", "Y"],
).astype(dtype)
values_2 = "111"
targ_blank = xr.DataArray(
[["a111", "bb111", "cccc111"], ["ddddd111", "eeee111", "fff111"]],
dims=["X", "Y"],
).astype(dtype)
targ_space = xr.DataArray(
[["a 111", "bb 111", "cccc 111"], ["ddddd 111", "eeee 111", "fff 111"]],
dims=["X", "Y"],
).astype(dtype)
targ_bars = xr.DataArray(
[["a||111", "bb||111", "cccc||111"], ["ddddd||111", "eeee||111", "fff||111"]],
dims=["X", "Y"],
).astype(dtype)
targ_comma = xr.DataArray(
[["a, 111", "bb, 111", "cccc, 111"], ["ddddd, 111", "eeee, 111", "fff, 111"]],
dims=["X", "Y"],
).astype(dtype)
res_blank = values_1.str.cat(values_2)
res_add = values_1.str + values_2
res_space = values_1.str.cat(values_2, sep=" ")
res_bars = values_1.str.cat(values_2, sep="||")
res_comma = values_1.str.cat(values_2, sep=", ")
assert res_blank.dtype == targ_blank.dtype
assert res_add.dtype == targ_blank.dtype
assert res_space.dtype == targ_space.dtype
assert res_bars.dtype == targ_bars.dtype
assert res_comma.dtype == targ_comma.dtype
assert_equal(res_blank, targ_blank)
assert_equal(res_add, targ_blank)
assert_equal(res_space, targ_space)
assert_equal(res_bars, targ_bars)
assert_equal(res_comma, targ_comma)
def test_cat_uniform(dtype) -> None:
values_1 = xr.DataArray(
[["a", "bb", "cccc"], ["ddddd", "eeee", "fff"]],
dims=["X", "Y"],
).astype(dtype)
values_2 = xr.DataArray(
[["11111", "222", "33"], ["4", "5555", "66"]],
dims=["X", "Y"],
)
targ_blank = xr.DataArray(
[["a11111", "bb222", "cccc33"], ["ddddd4", "eeee5555", "fff66"]],
dims=["X", "Y"],
).astype(dtype)
targ_space = xr.DataArray(
[["a 11111", "bb 222", "cccc 33"], ["ddddd 4", "eeee 5555", "fff 66"]],
dims=["X", "Y"],
).astype(dtype)
targ_bars = xr.DataArray(
[["a||11111", "bb||222", "cccc||33"], ["ddddd||4", "eeee||5555", "fff||66"]],
dims=["X", "Y"],
).astype(dtype)
targ_comma = xr.DataArray(
[["a, 11111", "bb, 222", "cccc, 33"], ["ddddd, 4", "eeee, 5555", "fff, 66"]],
dims=["X", "Y"],
).astype(dtype)
res_blank = values_1.str.cat(values_2)
res_add = values_1.str + values_2
res_space = values_1.str.cat(values_2, sep=" ")
res_bars = values_1.str.cat(values_2, sep="||")
res_comma = values_1.str.cat(values_2, sep=", ")
assert res_blank.dtype == targ_blank.dtype
assert res_add.dtype == targ_blank.dtype
assert res_space.dtype == targ_space.dtype
assert res_bars.dtype == targ_bars.dtype
assert res_comma.dtype == targ_comma.dtype
assert_equal(res_blank, targ_blank)
assert_equal(res_add, targ_blank)
assert_equal(res_space, targ_space)
assert_equal(res_bars, targ_bars)
assert_equal(res_comma, targ_comma)
def test_cat_broadcast_right(dtype) -> None:
values_1 = xr.DataArray(
[["a", "bb", "cccc"], ["ddddd", "eeee", "fff"]],
dims=["X", "Y"],
).astype(dtype)
values_2 = xr.DataArray(
["11111", "222", "33"],
dims=["Y"],
)
targ_blank = xr.DataArray(
[["a11111", "bb222", "cccc33"], ["ddddd11111", "eeee222", "fff33"]],
dims=["X", "Y"],
).astype(dtype)
targ_space = xr.DataArray(
[["a 11111", "bb 222", "cccc 33"], ["ddddd 11111", "eeee 222", "fff 33"]],
dims=["X", "Y"],
).astype(dtype)
targ_bars = xr.DataArray(
[["a||11111", "bb||222", "cccc||33"], ["ddddd||11111", "eeee||222", "fff||33"]],
dims=["X", "Y"],
).astype(dtype)
targ_comma = xr.DataArray(
[["a, 11111", "bb, 222", "cccc, 33"], ["ddddd, 11111", "eeee, 222", "fff, 33"]],
dims=["X", "Y"],
).astype(dtype)
res_blank = values_1.str.cat(values_2)
res_add = values_1.str + values_2
res_space = values_1.str.cat(values_2, sep=" ")
res_bars = values_1.str.cat(values_2, sep="||")
res_comma = values_1.str.cat(values_2, sep=", ")
assert res_blank.dtype == targ_blank.dtype
assert res_add.dtype == targ_blank.dtype
assert res_space.dtype == targ_space.dtype
assert res_bars.dtype == targ_bars.dtype
assert res_comma.dtype == targ_comma.dtype
assert_equal(res_blank, targ_blank)
assert_equal(res_add, targ_blank)
assert_equal(res_space, targ_space)
assert_equal(res_bars, targ_bars)
assert_equal(res_comma, targ_comma)
def test_cat_broadcast_left(dtype) -> None:
values_1 = xr.DataArray(
["a", "bb", "cccc"],
dims=["Y"],
).astype(dtype)
values_2 = xr.DataArray(
[["11111", "222", "33"], ["4", "5555", "66"]],
dims=["X", "Y"],
)
targ_blank = (
xr.DataArray(
[["a11111", "bb222", "cccc33"], ["a4", "bb5555", "cccc66"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
targ_space = (
xr.DataArray(
[["a 11111", "bb 222", "cccc 33"], ["a 4", "bb 5555", "cccc 66"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
targ_bars = (
xr.DataArray(
[["a||11111", "bb||222", "cccc||33"], ["a||4", "bb||5555", "cccc||66"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
targ_comma = (
xr.DataArray(
[["a, 11111", "bb, 222", "cccc, 33"], ["a, 4", "bb, 5555", "cccc, 66"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
res_blank = values_1.str.cat(values_2)
res_add = values_1.str + values_2
res_space = values_1.str.cat(values_2, sep=" ")
res_bars = values_1.str.cat(values_2, sep="||")
res_comma = values_1.str.cat(values_2, sep=", ")
assert res_blank.dtype == targ_blank.dtype
assert res_add.dtype == targ_blank.dtype
assert res_space.dtype == targ_space.dtype
assert res_bars.dtype == targ_bars.dtype
assert res_comma.dtype == targ_comma.dtype
assert_equal(res_blank, targ_blank)
assert_equal(res_add, targ_blank)
assert_equal(res_space, targ_space)
assert_equal(res_bars, targ_bars)
assert_equal(res_comma, targ_comma)
def test_cat_broadcast_both(dtype) -> None:
values_1 = xr.DataArray(
["a", "bb", "cccc"],
dims=["Y"],
).astype(dtype)
values_2 = xr.DataArray(
["11111", "4"],
dims=["X"],
)
targ_blank = (
xr.DataArray(
[["a11111", "bb11111", "cccc11111"], ["a4", "bb4", "cccc4"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
targ_space = (
xr.DataArray(
[["a 11111", "bb 11111", "cccc 11111"], ["a 4", "bb 4", "cccc 4"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
targ_bars = (
xr.DataArray(
[["a||11111", "bb||11111", "cccc||11111"], ["a||4", "bb||4", "cccc||4"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
targ_comma = (
xr.DataArray(
[["a, 11111", "bb, 11111", "cccc, 11111"], ["a, 4", "bb, 4", "cccc, 4"]],
dims=["X", "Y"],
)
.astype(dtype)
.T
)
res_blank = values_1.str.cat(values_2)
res_add = values_1.str + values_2
res_space = values_1.str.cat(values_2, sep=" ")
res_bars = values_1.str.cat(values_2, sep="||")
res_comma = values_1.str.cat(values_2, sep=", ")
assert res_blank.dtype == targ_blank.dtype
assert res_add.dtype == targ_blank.dtype
assert res_space.dtype == targ_space.dtype
assert res_bars.dtype == targ_bars.dtype
assert res_comma.dtype == targ_comma.dtype
assert_equal(res_blank, targ_blank)
assert_equal(res_add, targ_blank)
assert_equal(res_space, targ_space)
assert_equal(res_bars, targ_bars)
assert_equal(res_comma, targ_comma)
def test_cat_multi() -> None:
values_1 = xr.DataArray(
["11111", "4"],
dims=["X"],
)
values_2 = xr.DataArray(
["a", "bb", "cccc"],
dims=["Y"],
).astype(np.bytes_)
values_3 = np.array(3.4)
values_4 = ""
values_5 = np.array("", dtype=np.unicode_)
sep = xr.DataArray(
[" ", ", "],
dims=["ZZ"],
).astype(np.unicode_)
expected = xr.DataArray(
[
[
["11111 a 3.4 ", "11111, a, 3.4, , "],
["11111 bb 3.4 ", "11111, bb, 3.4, , "],
["11111 cccc 3.4 ", "11111, cccc, 3.4, , "],
],
[
["4 a 3.4 ", "4, a, 3.4, , "],
["4 bb 3.4 ", "4, bb, 3.4, , "],
["4 cccc 3.4 ", "4, cccc, 3.4, , "],
],
],
dims=["X", "Y", "ZZ"],
).astype(np.unicode_)
res = values_1.str.cat(values_2, values_3, values_4, values_5, sep=sep)
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_join_scalar(dtype) -> None:
values = xr.DataArray("aaa").astype(dtype)
targ = xr.DataArray("aaa").astype(dtype)
res_blank = values.str.join()
res_space = values.str.join(sep=" ")
assert res_blank.dtype == targ.dtype
assert res_space.dtype == targ.dtype
assert_identical(res_blank, targ)
assert_identical(res_space, targ)
def test_join_vector(dtype) -> None:
values = xr.DataArray(
["a", "bb", "cccc"],
dims=["Y"],
).astype(dtype)
targ_blank = xr.DataArray("abbcccc").astype(dtype)
targ_space = xr.DataArray("a bb cccc").astype(dtype)
res_blank_none = values.str.join()
res_blank_y = values.str.join(dim="Y")
res_space_none = values.str.join(sep=" ")
res_space_y = values.str.join(dim="Y", sep=" ")
assert res_blank_none.dtype == targ_blank.dtype
assert res_blank_y.dtype == targ_blank.dtype
assert res_space_none.dtype == targ_space.dtype
assert res_space_y.dtype == targ_space.dtype
assert_identical(res_blank_none, targ_blank)
assert_identical(res_blank_y, targ_blank)
assert_identical(res_space_none, targ_space)
assert_identical(res_space_y, targ_space)
def test_join_2d(dtype) -> None:
values = xr.DataArray(
[["a", "bb", "cccc"], ["ddddd", "eeee", "fff"]],
dims=["X", "Y"],
).astype(dtype)
targ_blank_x = xr.DataArray(
["addddd", "bbeeee", "ccccfff"],
dims=["Y"],
).astype(dtype)
targ_space_x = xr.DataArray(
["a ddddd", "bb eeee", "cccc fff"],
dims=["Y"],
).astype(dtype)
targ_blank_y = xr.DataArray(
["abbcccc", "dddddeeeefff"],
dims=["X"],
).astype(dtype)
targ_space_y = xr.DataArray(
["a bb cccc", "ddddd eeee fff"],
dims=["X"],
).astype(dtype)
res_blank_x = values.str.join(dim="X")
res_blank_y = values.str.join(dim="Y")
res_space_x = values.str.join(dim="X", sep=" ")
res_space_y = values.str.join(dim="Y", sep=" ")
assert res_blank_x.dtype == targ_blank_x.dtype
assert res_blank_y.dtype == targ_blank_y.dtype
assert res_space_x.dtype == targ_space_x.dtype
assert res_space_y.dtype == targ_space_y.dtype
assert_identical(res_blank_x, targ_blank_x)
assert_identical(res_blank_y, targ_blank_y)
assert_identical(res_space_x, targ_space_x)
assert_identical(res_space_y, targ_space_y)
with pytest.raises(
ValueError, match="Dimension must be specified for multidimensional arrays."
):
values.str.join()
def test_join_broadcast(dtype) -> None:
values = xr.DataArray(
["a", "bb", "cccc"],
dims=["X"],
).astype(dtype)
sep = xr.DataArray(
[" ", ", "],
dims=["ZZ"],
).astype(dtype)
expected = xr.DataArray(
["a bb cccc", "a, bb, cccc"],
dims=["ZZ"],
).astype(dtype)
res = values.str.join(sep=sep)
assert res.dtype == expected.dtype
assert_identical(res, expected)
def test_format_scalar() -> None:
values = xr.DataArray(
["{}.{Y}.{ZZ}", "{},{},{X},{X}", "{X}-{Y}-{ZZ}"],
dims=["X"],
).astype(np.unicode_)
pos0 = 1
pos1 = 1.2
pos2 = "2.3"
X = "'test'"
Y = "X"
ZZ = None
W = "NO!"
expected = xr.DataArray(
["1.X.None", "1,1.2,'test','test'", "'test'-X-None"],
dims=["X"],
).astype(np.unicode_)
res = values.str.format(pos0, pos1, pos2, X=X, Y=Y, ZZ=ZZ, W=W)
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_format_broadcast() -> None:
values = xr.DataArray(
["{}.{Y}.{ZZ}", "{},{},{X},{X}", "{X}-{Y}-{ZZ}"],
dims=["X"],
).astype(np.unicode_)
pos0 = 1
pos1 = 1.2
pos2 = xr.DataArray(
["2.3", "3.44444"],
dims=["YY"],
)
X = "'test'"
Y = "X"
ZZ = None
W = "NO!"
expected = xr.DataArray(
[
["1.X.None", "1.X.None"],
["1,1.2,'test','test'", "1,1.2,'test','test'"],
["'test'-X-None", "'test'-X-None"],
],
dims=["X", "YY"],
).astype(np.unicode_)
res = values.str.format(pos0, pos1, pos2, X=X, Y=Y, ZZ=ZZ, W=W)
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_mod_scalar() -> None:
values = xr.DataArray(
["%s.%s.%s", "%s,%s,%s", "%s-%s-%s"],
dims=["X"],
).astype(np.unicode_)
pos0 = 1
pos1 = 1.2
pos2 = "2.3"
expected = xr.DataArray(
["1.1.2.2.3", "1,1.2,2.3", "1-1.2-2.3"],
dims=["X"],
).astype(np.unicode_)
res = values.str % (pos0, pos1, pos2)
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_mod_dict() -> None:
values = xr.DataArray(
["%(a)s.%(a)s.%(b)s", "%(b)s,%(c)s,%(b)s", "%(c)s-%(b)s-%(a)s"],
dims=["X"],
).astype(np.unicode_)
a = 1
b = 1.2
c = "2.3"
expected = xr.DataArray(
["1.1.1.2", "1.2,2.3,1.2", "2.3-1.2-1"],
dims=["X"],
).astype(np.unicode_)
res = values.str % {"a": a, "b": b, "c": c}
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_mod_broadcast_single() -> None:
values = xr.DataArray(
["%s_1", "%s_2", "%s_3"],
dims=["X"],
).astype(np.unicode_)
pos = xr.DataArray(
["2.3", "3.44444"],
dims=["YY"],
)
expected = xr.DataArray(
[["2.3_1", "3.44444_1"], ["2.3_2", "3.44444_2"], ["2.3_3", "3.44444_3"]],
dims=["X", "YY"],
).astype(np.unicode_)
res = values.str % pos
assert res.dtype == expected.dtype
assert_equal(res, expected)
def test_mod_broadcast_multi() -> None:
values = xr.DataArray(
["%s.%s.%s", "%s,%s,%s", "%s-%s-%s"],
dims=["X"],
).astype(np.unicode_)
pos0 = 1
pos1 = 1.2
pos2 = xr.DataArray(
["2.3", "3.44444"],
dims=["YY"],
)
expected = xr.DataArray(
[
["1.1.2.2.3", "1.1.2.3.44444"],
["1,1.2,2.3", "1,1.2,3.44444"],
["1-1.2-2.3", "1-1.2-3.44444"],
],
dims=["X", "YY"],
).astype(np.unicode_)
res = values.str % (pos0, pos1, pos2)
assert res.dtype == expected.dtype
assert_equal(res, expected)