# The StringAccessor class defined below is an adaptation of the # pandas string methods source code (see pd.core.strings) # 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. import codecs import re import textwrap import numpy as np from .computation import apply_ufunc _cpython_optimized_encoders = ( "utf-8", "utf8", "latin-1", "latin1", "iso-8859-1", "mbcs", "ascii", ) _cpython_optimized_decoders = _cpython_optimized_encoders + ("utf-16", "utf-32") def _is_str_like(x): return isinstance(x, str) or isinstance(x, bytes) class StringAccessor: """Vectorized string functions for string-like arrays. Similar to pandas, fields can be accessed through the `.str` attribute for applicable DataArrays. >>> da = xr.DataArray(["some", "text", "in", "an", "array"]) >>> da.str.len() array([4, 4, 2, 2, 5]) Dimensions without coordinates: dim_0 """ __slots__ = ("_obj",) def __init__(self, obj): self._obj = obj def _apply(self, f, dtype=None): # TODO handling of na values ? if dtype is None: dtype = self._obj.dtype g = np.vectorize(f, otypes=[dtype]) return apply_ufunc(g, self._obj, dask="parallelized", output_dtypes=[dtype]) def len(self): """ Compute the length of each string in the array. Returns ------- lengths array : array of int """ return self._apply(len, dtype=int) def __getitem__(self, key): if isinstance(key, slice): return self.slice(start=key.start, stop=key.stop, step=key.step) else: return self.get(key) def get(self, i, default=""): """ Extract character number `i` from each string in the array. Parameters ---------- i : int Position of element to extract. default : optional Value for out-of-range index. If not specified (None) defaults to an empty string. Returns ------- items : array of object """ s = slice(-1, None) if i == -1 else slice(i, i + 1) def f(x): item = x[s] return item if item else default return self._apply(f) def slice(self, start=None, stop=None, step=None): """ Slice substrings from each string in the array. Parameters ---------- start : int, optional Start position for slice operation. stop : int, optional Stop position for slice operation. step : int, optional Step size for slice operation. Returns ------- sliced strings : same type as values """ s = slice(start, stop, step) f = lambda x: x[s] return self._apply(f) def slice_replace(self, start=None, stop=None, repl=""): """ Replace a positional slice of a string with another value. Parameters ---------- start : int, optional Left index position to use for the slice. If not specified (None), the slice is unbounded on the left, i.e. slice from the start of the string. stop : int, optional Right index position to use for the slice. If not specified (None), the slice is unbounded on the right, i.e. slice until the end of the string. repl : str, optional String for replacement. If not specified, the sliced region is replaced with an empty string. Returns ------- replaced : same type as values """ repl = self._obj.dtype.type(repl) def f(x): if len(x[start:stop]) == 0: local_stop = start else: local_stop = stop y = self._obj.dtype.type("") if start is not None: y += x[:start] y += repl if stop is not None: y += x[local_stop:] return y return self._apply(f) def capitalize(self): """ Convert strings in the array to be capitalized. Returns ------- capitalized : same type as values """ return self._apply(lambda x: x.capitalize()) def lower(self): """ Convert strings in the array to lowercase. Returns ------- lowerd : same type as values """ return self._apply(lambda x: x.lower()) def swapcase(self): """ Convert strings in the array to be swapcased. Returns ------- swapcased : same type as values """ return self._apply(lambda x: x.swapcase()) def title(self): """ Convert strings in the array to titlecase. Returns ------- titled : same type as values """ return self._apply(lambda x: x.title()) def upper(self): """ Convert strings in the array to uppercase. Returns ------- uppered : same type as values """ return self._apply(lambda x: x.upper()) def isalnum(self): """ Check whether all characters in each string are alphanumeric. Returns ------- isalnum : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isalnum(), dtype=bool) def isalpha(self): """ Check whether all characters in each string are alphabetic. Returns ------- isalpha : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isalpha(), dtype=bool) def isdecimal(self): """ Check whether all characters in each string are decimal. Returns ------- isdecimal : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isdecimal(), dtype=bool) def isdigit(self): """ Check whether all characters in each string are digits. Returns ------- isdigit : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isdigit(), dtype=bool) def islower(self): """ Check whether all characters in each string are lowercase. Returns ------- islower : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.islower(), dtype=bool) def isnumeric(self): """ Check whether all characters in each string are numeric. Returns ------- isnumeric : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isnumeric(), dtype=bool) def isspace(self): """ Check whether all characters in each string are spaces. Returns ------- isspace : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isspace(), dtype=bool) def istitle(self): """ Check whether all characters in each string are titlecase. Returns ------- istitle : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.istitle(), dtype=bool) def isupper(self): """ Check whether all characters in each string are uppercase. Returns ------- isupper : array of bool Array of boolean values with the same shape as the original array. """ return self._apply(lambda x: x.isupper(), dtype=bool) def count(self, pat, flags=0): """ Count occurrences of pattern in each string of the array. This function is used to count the number of times a particular regex pattern is repeated in each of the string elements of the :class:`~xarray.DataArray`. Parameters ---------- pat : str Valid regular expression. flags : int, default: 0 Flags for the `re` module. Use 0 for no flags. For a complete list, `see here `_. Returns ------- counts : array of int """ pat = self._obj.dtype.type(pat) regex = re.compile(pat, flags=flags) f = lambda x: len(regex.findall(x)) return self._apply(f, dtype=int) def startswith(self, pat): """ Test if the start of each string in the array matches a pattern. Parameters ---------- pat : str Character sequence. Regular expressions are not accepted. Returns ------- startswith : array of bool An array of booleans indicating whether the given pattern matches the start of each string element. """ pat = self._obj.dtype.type(pat) f = lambda x: x.startswith(pat) return self._apply(f, dtype=bool) def endswith(self, pat): """ Test if the end of each string in the array matches a pattern. Parameters ---------- pat : str Character sequence. Regular expressions are not accepted. Returns ------- endswith : array of bool A Series of booleans indicating whether the given pattern matches the end of each string element. """ pat = self._obj.dtype.type(pat) f = lambda x: x.endswith(pat) return self._apply(f, dtype=bool) def pad(self, width, side="left", fillchar=" "): """ Pad strings in the array up to width. Parameters ---------- width : int Minimum width of resulting string; additional characters will be filled with character defined in `fillchar`. side : {"left", "right", "both"}, default: "left" Side from which to fill resulting string. fillchar : str, default: " " Additional character for filling, default is whitespace. Returns ------- filled : same type as values Array with a minimum number of char in each element. """ width = int(width) fillchar = self._obj.dtype.type(fillchar) if len(fillchar) != 1: raise TypeError("fillchar must be a character, not str") if side == "left": f = lambda s: s.rjust(width, fillchar) elif side == "right": f = lambda s: s.ljust(width, fillchar) elif side == "both": f = lambda s: s.center(width, fillchar) else: # pragma: no cover raise ValueError("Invalid side") return self._apply(f) def center(self, width, fillchar=" "): """ Pad left and right side of each string in the array. Parameters ---------- width : int Minimum width of resulting string; additional characters will be filled with ``fillchar`` fillchar : str, default: " " Additional character for filling, default is whitespace Returns ------- filled : same type as values """ return self.pad(width, side="both", fillchar=fillchar) def ljust(self, width, fillchar=" "): """ Pad right side of each string in the array. Parameters ---------- width : int Minimum width of resulting string; additional characters will be filled with ``fillchar`` fillchar : str, default: " " Additional character for filling, default is whitespace Returns ------- filled : same type as values """ return self.pad(width, side="right", fillchar=fillchar) def rjust(self, width, fillchar=" "): """ Pad left side of each string in the array. Parameters ---------- width : int Minimum width of resulting string; additional characters will be filled with ``fillchar`` fillchar : str, default: " " Additional character for filling, default is whitespace Returns ------- filled : same type as values """ return self.pad(width, side="left", fillchar=fillchar) def zfill(self, width): """ Pad each string in the array by prepending '0' characters. Strings in the array are padded with '0' characters on the left of the string to reach a total string length `width`. Strings in the array with length greater or equal to `width` are unchanged. Parameters ---------- width : int Minimum length of resulting string; strings with length less than `width` be prepended with '0' characters. Returns ------- filled : same type as values """ return self.pad(width, side="left", fillchar="0") def contains(self, pat, case=True, flags=0, regex=True): """ Test if pattern or regex is contained within each string of the array. Return boolean array based on whether a given pattern or regex is contained within a string of the array. Parameters ---------- pat : str Character sequence or regular expression. case : bool, default: True If True, case sensitive. flags : int, default: 0 Flags to pass through to the re module, e.g. re.IGNORECASE. ``0`` means no flags. regex : bool, default: True If True, assumes the pat is a regular expression. If False, treats the pat as a literal string. Returns ------- contains : array of bool An array of boolean values indicating whether the given pattern is contained within the string of each element of the array. """ pat = self._obj.dtype.type(pat) if regex: if not case: flags |= re.IGNORECASE regex = re.compile(pat, flags=flags) if regex.groups > 0: # pragma: no cover raise ValueError("This pattern has match groups.") f = lambda x: bool(regex.search(x)) else: if case: f = lambda x: pat in x else: uppered = self._obj.str.upper() return uppered.str.contains(pat.upper(), regex=False) return self._apply(f, dtype=bool) def match(self, pat, case=True, flags=0): """ Determine if each string in the array matches a regular expression. Parameters ---------- pat : str Character sequence or regular expression case : bool, default: True If True, case sensitive flags : int, default: 0 re module flags, e.g. re.IGNORECASE. ``0`` means no flags Returns ------- matched : array of bool """ if not case: flags |= re.IGNORECASE pat = self._obj.dtype.type(pat) regex = re.compile(pat, flags=flags) f = lambda x: bool(regex.match(x)) return self._apply(f, dtype=bool) def strip(self, to_strip=None, side="both"): """ Remove leading and trailing characters. Strip whitespaces (including newlines) or a set of specified characters from each string in the array from left and/or right sides. Parameters ---------- to_strip : str or None, default: None Specifying the set of characters to be removed. All combinations of this set of characters will be stripped. If None then whitespaces are removed. side : {"left", "right", "both"}, default: "left" Side from which to strip. Returns ------- stripped : same type as values """ if to_strip is not None: to_strip = self._obj.dtype.type(to_strip) if side == "both": f = lambda x: x.strip(to_strip) elif side == "left": f = lambda x: x.lstrip(to_strip) elif side == "right": f = lambda x: x.rstrip(to_strip) else: # pragma: no cover raise ValueError("Invalid side") return self._apply(f) def lstrip(self, to_strip=None): """ Remove leading characters. Strip whitespaces (including newlines) or a set of specified characters from each string in the array from the left side. Parameters ---------- to_strip : str or None, default: None Specifying the set of characters to be removed. All combinations of this set of characters will be stripped. If None then whitespaces are removed. Returns ------- stripped : same type as values """ return self.strip(to_strip, side="left") def rstrip(self, to_strip=None): """ Remove trailing characters. Strip whitespaces (including newlines) or a set of specified characters from each string in the array from the right side. Parameters ---------- to_strip : str or None, default: None Specifying the set of characters to be removed. All combinations of this set of characters will be stripped. If None then whitespaces are removed. Returns ------- stripped : same type as values """ return self.strip(to_strip, side="right") def wrap(self, width, **kwargs): """ Wrap long strings in the array in paragraphs with length less than `width`. This method has the same keyword parameters and defaults as :class:`textwrap.TextWrapper`. Parameters ---------- width : int Maximum line-width **kwargs keyword arguments passed into :class:`textwrap.TextWrapper`. Returns ------- wrapped : same type as values """ tw = textwrap.TextWrapper(width=width, **kwargs) f = lambda x: "\n".join(tw.wrap(x)) return self._apply(f) def translate(self, table): """ Map characters of each string through the given mapping table. Parameters ---------- table : dict A a mapping of Unicode ordinals to Unicode ordinals, strings, or None. Unmapped characters are left untouched. Characters mapped to None are deleted. :meth:`str.maketrans` is a helper function for making translation tables. Returns ------- translated : same type as values """ f = lambda x: x.translate(table) return self._apply(f) def repeat(self, repeats): """ Duplicate each string in the array. Parameters ---------- repeats : int Number of repetitions. Returns ------- repeated : same type as values Array of repeated string objects. """ f = lambda x: repeats * x return self._apply(f) def find(self, sub, start=0, end=None, side="left"): """ Return lowest or highest indexes in each strings in the array where the substring is fully contained between [start:end]. Return -1 on failure. Parameters ---------- sub : str Substring being searched start : int Left edge index end : int Right edge index side : {"left", "right"}, default: "left" Starting side for search. Returns ------- found : array of int """ sub = self._obj.dtype.type(sub) if side == "left": method = "find" elif side == "right": method = "rfind" else: # pragma: no cover raise ValueError("Invalid side") if end is None: f = lambda x: getattr(x, method)(sub, start) else: f = lambda x: getattr(x, method)(sub, start, end) return self._apply(f, dtype=int) def rfind(self, sub, start=0, end=None): """ Return highest indexes in each strings in the array where the substring is fully contained between [start:end]. Return -1 on failure. Parameters ---------- sub : str Substring being searched start : int Left edge index end : int Right edge index Returns ------- found : array of int """ return self.find(sub, start=start, end=end, side="right") def index(self, sub, start=0, end=None, side="left"): """ Return lowest or highest indexes in each strings where the substring is fully contained between [start:end]. This is the same as ``str.find`` except instead of returning -1, it raises a ValueError when the substring is not found. Parameters ---------- sub : str Substring being searched start : int Left edge index end : int Right edge index side : {"left", "right"}, default: "left" Starting side for search. Returns ------- found : array of int """ sub = self._obj.dtype.type(sub) if side == "left": method = "index" elif side == "right": method = "rindex" else: # pragma: no cover raise ValueError("Invalid side") if end is None: f = lambda x: getattr(x, method)(sub, start) else: f = lambda x: getattr(x, method)(sub, start, end) return self._apply(f, dtype=int) def rindex(self, sub, start=0, end=None): """ Return highest indexes in each strings where the substring is fully contained between [start:end]. This is the same as ``str.rfind`` except instead of returning -1, it raises a ValueError when the substring is not found. Parameters ---------- sub : str Substring being searched start : int Left edge index end : int Right edge index Returns ------- found : array of int """ return self.index(sub, start=start, end=end, side="right") def replace(self, pat, repl, n=-1, case=None, flags=0, regex=True): """ Replace occurrences of pattern/regex in the array with some string. Parameters ---------- pat : str or re.Pattern String can be a character sequence or regular expression. repl : str or callable Replacement string or a callable. The callable is passed the regex match object and must return a replacement string to be used. See :func:`re.sub`. n : int, default: -1 Number of replacements to make from start. Use ``-1`` to replace all. case : bool, default: None - If True, case sensitive (the default if `pat` is a string) - Set to False for case insensitive - Cannot be set if `pat` is a compiled regex flags : int, default: 0 - re module flags, e.g. re.IGNORECASE. Use ``0`` for no flags. - Cannot be set if `pat` is a compiled regex regex : bool, default: True - If True, assumes the passed-in pattern is a regular expression. - If False, treats the pattern as a literal string - Cannot be set to False if `pat` is a compiled regex or `repl` is a callable. Returns ------- replaced : same type as values A copy of the object with all matching occurrences of `pat` replaced by `repl`. """ if not (_is_str_like(repl) or callable(repl)): # pragma: no cover raise TypeError("repl must be a string or callable") if _is_str_like(pat): pat = self._obj.dtype.type(pat) if _is_str_like(repl): repl = self._obj.dtype.type(repl) is_compiled_re = isinstance(pat, type(re.compile(""))) if regex: if is_compiled_re: if (case is not None) or (flags != 0): raise ValueError( "case and flags cannot be set when pat is a compiled regex" ) else: # not a compiled regex # set default case if case is None: case = True # add case flag, if provided if case is False: flags |= re.IGNORECASE if is_compiled_re or len(pat) > 1 or flags or callable(repl): n = n if n >= 0 else 0 compiled = re.compile(pat, flags=flags) f = lambda x: compiled.sub(repl=repl, string=x, count=n) else: f = lambda x: x.replace(pat, repl, n) else: if is_compiled_re: raise ValueError( "Cannot use a compiled regex as replacement " "pattern with regex=False" ) if callable(repl): raise ValueError("Cannot use a callable replacement when regex=False") f = lambda x: x.replace(pat, repl, n) return self._apply(f) def decode(self, encoding, errors="strict"): """ Decode character string in the array using indicated encoding. Parameters ---------- encoding : str errors : str, optional Returns ------- decoded : same type as values """ if encoding in _cpython_optimized_decoders: f = lambda x: x.decode(encoding, errors) else: decoder = codecs.getdecoder(encoding) f = lambda x: decoder(x, errors)[0] return self._apply(f, dtype=np.str_) def encode(self, encoding, errors="strict"): """ Encode character string in the array using indicated encoding. Parameters ---------- encoding : str errors : str, optional Returns ------- encoded : same type as values """ if encoding in _cpython_optimized_encoders: f = lambda x: x.encode(encoding, errors) else: encoder = codecs.getencoder(encoding) f = lambda x: encoder(x, errors)[0] return self._apply(f, dtype=np.bytes_)