strings.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import re
from typing import cast, Literal, Optional, Union

import bigframes_vendored.constants as constants
import bigframes_vendored.pandas.core.strings.accessor as vendorstr

from bigframes.core import log_adapter
import bigframes.dataframe as df
import bigframes.operations as ops
from bigframes.operations._op_converters import convert_index, convert_slice
import bigframes.operations.base
import bigframes.series as series

# Maps from python to re2
REGEXP_FLAGS = {
    re.IGNORECASE: "i",
    re.MULTILINE: "m",
    re.DOTALL: "s",
}


@log_adapter.class_logger
class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMethods):
    __doc__ = vendorstr.StringMethods.__doc__

    def __getitem__(self, key: Union[int, slice]) -> series.Series:
        if isinstance(key, int):
            return self._apply_unary_op(convert_index(key))
        elif isinstance(key, slice):
            return self._apply_unary_op(convert_slice(key))
        else:
            raise ValueError(f"key must be an int or slice, got {type(key).__name__}")

    def find(
        self,
        sub: str,
        start: Optional[int] = None,
        end: Optional[int] = None,
    ) -> series.Series:
        return self._apply_unary_op(ops.StrFindOp(substr=sub, start=start, end=end))

    def len(self) -> series.Series:
        return self._apply_unary_op(ops.len_op)

    def lower(self) -> series.Series:
        return self._apply_unary_op(ops.lower_op)

    def reverse(self) -> series.Series:
        """Reverse strings in the Series.

        **Examples:**

            >>> import bigframes.pandas as bpd
            >>> bpd.options.display.progress_bar = None

            >>> s = bpd.Series(["apple", "banana", "", bpd.NA])
            >>> s.str.reverse()
            0     elppa
            1    ananab
            2
            3      <NA>
            dtype: string

        Returns:
            bigframes.series.Series: A Series of booleans indicating whether the given
                pattern matches the start of each string element.
        """
        # reverse method is in ibis, not pandas.
        return self._apply_unary_op(ops.reverse_op)

    def slice(
        self,
        start: Optional[int] = None,
        stop: Optional[int] = None,
    ) -> series.Series:
        return self._apply_unary_op(ops.StrSliceOp(start=start, end=stop))

    def strip(self) -> series.Series:
        return self._apply_unary_op(ops.strip_op)

    def upper(self) -> series.Series:
        return self._apply_unary_op(ops.upper_op)

    def isnumeric(self) -> series.Series:
        return self._apply_unary_op(ops.isnumeric_op)

    def isalpha(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.isalpha_op)

    def isdigit(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.isdigit_op)

    def isdecimal(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.isdecimal_op)

    def isalnum(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.isalnum_op)

    def isspace(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.isspace_op)

    def islower(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.islower_op)

    def isupper(
        self,
    ) -> series.Series:
        return self._apply_unary_op(ops.isupper_op)

    def rstrip(self) -> series.Series:
        return self._apply_unary_op(ops.rstrip_op)

    def lstrip(self) -> series.Series:
        return self._apply_unary_op(ops.lstrip_op)

    def repeat(self, repeats: int) -> series.Series:
        return self._apply_unary_op(ops.StrRepeatOp(repeats=repeats))

    def capitalize(self) -> series.Series:
        return self._apply_unary_op(ops.capitalize_op)

    def match(self, pat, case=True, flags=0) -> series.Series:
        # \A anchors start of entire string rather than start of any line in multiline mode
        adj_pat = rf"\A{pat}"
        return self.contains(pat=adj_pat, case=case, flags=flags)

    def fullmatch(self, pat, case=True, flags=0) -> series.Series:
        # \A anchors start of entire string rather than start of any line in multiline mode
        # \z likewise anchors to the end of the entire multiline string
        adj_pat = rf"\A{pat}\z"
        return self.contains(pat=adj_pat, case=case, flags=flags)

    def get(self, i: int) -> series.Series:
        return self._apply_unary_op(ops.StrGetOp(i=i))

    def pad(self, width, side="left", fillchar=" ") -> series.Series:
        return self._apply_unary_op(
            ops.StrPadOp(length=width, fillchar=fillchar, side=side)
        )

    def ljust(self, width, fillchar=" ") -> series.Series:
        return self._apply_unary_op(
            ops.StrPadOp(length=width, fillchar=fillchar, side="right")
        )

    def rjust(self, width, fillchar=" ") -> series.Series:
        return self._apply_unary_op(
            ops.StrPadOp(length=width, fillchar=fillchar, side="left")
        )

    def contains(
        self, pat, case: bool = True, flags: int = 0, *, regex: bool = True
    ) -> series.Series:
        if not case:
            return self.contains(pat=pat, flags=flags | re.IGNORECASE, regex=True)
        if regex:
            re2flags = _parse_flags(flags)
            if re2flags:
                pat = re2flags + pat
            return self._apply_unary_op(ops.StrContainsRegexOp(pat=pat))
        else:
            return self._apply_unary_op(ops.StrContainsOp(pat=pat))

    def extract(self, pat: str, flags: int = 0) -> df.DataFrame:
        re2flags = _parse_flags(flags)
        if re2flags:
            pat = re2flags + pat
        compiled = re.compile(pat)
        if compiled.groups == 0:
            raise ValueError("No capture groups in 'pat'")

        results: list[str] = []
        block = self._block
        for i in range(compiled.groups):
            labels = [
                label
                for label, groupn in compiled.groupindex.items()
                if i + 1 == groupn
            ]
            label = labels[0] if labels else str(i)
            block, id = block.apply_unary_op(
                self._value_column,
                ops.StrExtractOp(pat=pat, n=i + 1),
                result_label=label,
            )
            results.append(id)
        block = block.select_columns(results)
        return df.DataFrame(block)

    def replace(
        self,
        pat: Union[str, re.Pattern],
        repl: str,
        *,
        case: Optional[bool] = None,
        flags: int = 0,
        regex: bool = False,
    ) -> series.Series:
        is_compiled = isinstance(pat, re.Pattern)
        patstr = cast(str, pat.pattern if is_compiled else pat)  # type: ignore
        if case is False:
            return self.replace(pat, repl, flags=flags | re.IGNORECASE, regex=True)
        if regex:
            re2flags = _parse_flags(flags)
            if re2flags:
                patstr = re2flags + patstr
            return self._apply_unary_op(ops.RegexReplaceStrOp(pat=patstr, repl=repl))
        else:
            if is_compiled:
                raise ValueError(
                    "Must set 'regex'=True if using compiled regex pattern."
                )
            return self._apply_unary_op(ops.ReplaceStrOp(pat=patstr, repl=repl))

    def startswith(
        self,
        pat: Union[str, tuple[str, ...]],
    ) -> series.Series:
        if not isinstance(pat, tuple):
            pat = (pat,)
        return self._apply_unary_op(ops.StartsWithOp(pat=pat))

    def endswith(
        self,
        pat: Union[str, tuple[str, ...]],
    ) -> series.Series:
        if not isinstance(pat, tuple):
            pat = (pat,)
        return self._apply_unary_op(ops.EndsWithOp(pat=pat))

    def split(
        self,
        pat: str = " ",
        regex: Union[bool, None] = None,
    ) -> series.Series:
        if regex is True or (regex is None and len(pat) > 1):
            raise NotImplementedError(
                "Regular expressions aren't currently supported. Please set "
                + f"`regex=False` and try again. {constants.FEEDBACK_LINK}"
            )
        return self._apply_unary_op(ops.StringSplitOp(pat=pat))

    def zfill(self, width: int) -> series.Series:
        return self._apply_unary_op(ops.ZfillOp(width=width))

    def center(self, width: int, fillchar: str = " ") -> series.Series:
        return self._apply_unary_op(
            ops.StrPadOp(length=width, fillchar=fillchar, side="both")
        )

    def cat(
        self,
        others: Union[str, series.Series],
        *,
        join: Literal["outer", "left"] = "left",
    ) -> series.Series:
        return self._apply_binary_op(others, ops.strconcat_op, alignment=join)

    def to_blob(self, connection: Optional[str] = None) -> series.Series:
        """Create a BigFrames Blob series from a series of URIs.

        .. note::
            BigFrames Blob is still under experiments. It may not work and subject to change in the future.


        Args:
            connection (str or None, default None):
                Connection to connect with remote service. str of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
                If None, use default connection in session context. BigQuery DataFrame will try to create the connection and attach
                permission if the connection isn't fully set up.

        Returns:
            bigframes.series.Series: Blob Series.

        """
        if not bigframes.options.experiments.blob:
            raise NotImplementedError()

        session = self._block.session
        connection = session._create_bq_connection(connection=connection)
        return self._apply_binary_op(connection, ops.obj_make_ref_op)


def _parse_flags(flags: int) -> Optional[str]:
    re2flags = []
    for reflag, re2flag in REGEXP_FLAGS.items():
        if flags & flags:
            re2flags.append(re2flag)
            flags = flags ^ reflag

    # Remaining flags couldn't be mapped to re2 engine
    if flags:
        raise NotImplementedError(
            f"Could not handle RegexFlag: {flags}. {constants.FEEDBACK_LINK}"
        )

    if re2flags:
        return "(?" + "".join(re2flags) + ")"
    else:
        return None