done

lib/python3.11/site-packages/narwhals/_sql/expr_str.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from typing import Any, Generic
+
+from narwhals._compliant import LazyExprNamespace
+from narwhals._compliant.any_namespace import StringNamespace
+from narwhals._sql.typing import SQLExprT
+
+
+class SQLExprStringNamespace(
+    LazyExprNamespace[SQLExprT], StringNamespace[SQLExprT], Generic[SQLExprT]
+):
+    def _lit(self, value: Any) -> SQLExprT:
+        return self.compliant._lit(value)  # type: ignore[no-any-return]
+
+    def _function(self, name: str, *args: Any) -> SQLExprT:
+        return self.compliant._function(name, *args)  # type: ignore[no-any-return]
+
+    def _when(self, condition: Any, value: Any, otherwise: Any | None = None) -> SQLExprT:
+        return self.compliant._when(condition, value, otherwise)  # type: ignore[no-any-return]
+
+    def contains(self, pattern: str, *, literal: bool) -> SQLExprT:
+        def func(expr: Any) -> Any:
+            if literal:
+                return self._function("contains", expr, self._lit(pattern))
+            return self._function("regexp_matches", expr, self._lit(pattern))
+
+        return self.compliant._with_elementwise(func)
+
+    def ends_with(self, suffix: str) -> SQLExprT:
+        return self.compliant._with_elementwise(
+            lambda expr: self._function("ends_with", expr, self._lit(suffix))
+        )
+
+    def len_chars(self) -> SQLExprT:
+        return self.compliant._with_elementwise(
+            lambda expr: self._function("length", expr)
+        )
+
+    def replace_all(
+        self, pattern: str, value: str | SQLExprT, *, literal: bool
+    ) -> SQLExprT:
+        fname: str = "replace" if literal else "regexp_replace"
+
+        options: list[Any] = []
+        if not literal and self.compliant._implementation.is_duckdb():
+            options = [self._lit("g")]
+
+        if isinstance(value, str):
+            return self.compliant._with_elementwise(
+                lambda expr: self._function(
+                    fname, expr, self._lit(pattern), self._lit(value), *options
+                )
+            )
+
+        return self.compliant._with_elementwise(
+            lambda expr, value: self._function(
+                fname, expr, self._lit(pattern), value, *options
+            ),
+            value=value,
+        )
+
+    def slice(self, offset: int, length: int | None) -> SQLExprT:
+        def func(expr: SQLExprT) -> SQLExprT:
+            col_length = self._function("length", expr)
+
+            _offset = (
+                col_length + self._lit(offset + 1)
+                if offset < 0
+                else self._lit(offset + 1)
+            )
+            _length = self._lit(length) if length is not None else col_length
+            return self._function("substr", expr, _offset, _length)
+
+        return self.compliant._with_elementwise(func)
+
+    def split(self, by: str) -> SQLExprT:
+        return self.compliant._with_elementwise(
+            lambda expr: self._function("str_split", expr, self._lit(by))
+        )
+
+    def starts_with(self, prefix: str) -> SQLExprT:
+        return self.compliant._with_elementwise(
+            lambda expr: self._function("starts_with", expr, self._lit(prefix))
+        )
+
+    def strip_chars(self, characters: str | None) -> SQLExprT:
+        import string
+
+        return self.compliant._with_elementwise(
+            lambda expr: self._function(
+                "trim",
+                expr,
+                self._lit(string.whitespace if characters is None else characters),
+            )
+        )
+
+    def to_lowercase(self) -> SQLExprT:
+        return self.compliant._with_elementwise(
+            lambda expr: self._function("lower", expr)
+        )
+
+    def to_uppercase(self) -> SQLExprT:
+        return self.compliant._with_elementwise(
+            lambda expr: self._function("upper", expr)
+        )
+
+    def zfill(self, width: int) -> SQLExprT:
+        # There is no built-in zfill function, so we need to implement it manually
+        # using string manipulation functions.
+
+        def func(expr: Any) -> Any:
+            less_than_width = self._function("length", expr) < self._lit(width)
+            zero, hyphen, plus = self._lit("0"), self._lit("-"), self._lit("+")
+
+            starts_with_minus = self._function("starts_with", expr, hyphen)
+            starts_with_plus = self._function("starts_with", expr, plus)
+            substring = self._function("substr", expr, self._lit(2))
+            padded_substring = self._function(
+                "lpad", substring, self._lit(width - 1), zero
+            )
+            return self._when(
+                starts_with_minus & less_than_width,
+                self._function("concat", hyphen, padded_substring),
+                self._when(
+                    starts_with_plus & less_than_width,
+                    self._function("concat", plus, padded_substring),
+                    self._when(
+                        less_than_width,
+                        self._function("lpad", expr, self._lit(width), zero),
+                        expr,
+                    ),
+                ),
+            )
+
+        # can't use `_with_elementwise` due to `when` operator.
+        # TODO(unassigned): implement `window_func` like we do in `Expr.cast`
+        return self.compliant._with_callable(func)