Repository: narwhals-dev/narwhals
Pull request #305 (2024-06-16 10:32:45)

enh: `slice`, `starts_with`, `tail` methods for `str` namespace
## What type of PR is this? (check all applicable)

- [x] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

None

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.

Implements `slice`, `starts_with` and `tail` methods in the str namespace for Expr and Series.

It is a quite large one, but slice implementation allowed to refactor head and tail, however using slice does not allow for negative indexes as polars does. Since they were not supported anyway in the first place, I moved to this implementation but happy to discuss how to work around it.



Patch:
diff --git a/docs/api-reference/expressions_str.md b/docs/api-reference/expressions_str.md
index 5d2f3deb..74e43bbd 100644
--- a/docs/api-reference/expressions_str.md
+++ b/docs/api-reference/expressions_str.md
@@ -4,8 +4,12 @@
     handler: python
     options:
       members:
+        - contains
         - ends_with
         - head
+        - slice
+        - starts_with
+        - tail
         - to_datetime
       show_source: false
       show_bases: false
diff --git a/docs/api-reference/series_str.md b/docs/api-reference/series_str.md
index 683075e6..3405b8a4 100644
--- a/docs/api-reference/series_str.md
+++ b/docs/api-reference/series_str.md
@@ -4,7 +4,11 @@
     handler: python
     options:
       members:
+        - contains
         - ends_with
         - head
+        - slice
+        - starts_with
+        - tail
       show_source: false
       show_bases: false
diff --git a/mkdocs.yml b/mkdocs.yml
index d1d4b128..64db7b7e 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -1,5 +1,7 @@
 site_name: Narwhals
 repo_url: https://github.com/narwhals-dev/narwhals.git
+watch:
+  - narwhals
 nav:
   - Home: index.md
   - Why: why.md
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 97c7f821..412d62b9 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -307,6 +307,14 @@ class PandasExprStringNamespace:
     def __init__(self, expr: PandasExpr) -> None:
         self._expr = expr
 
+    def starts_with(self, prefix: str) -> PandasExpr:
+        return reuse_series_namespace_implementation(
+            self._expr,
+            "str",
+            "starts_with",
+            prefix,
+        )
+
     def ends_with(self, suffix: str) -> PandasExpr:
         return reuse_series_namespace_implementation(
             self._expr,
@@ -324,12 +332,9 @@ class PandasExprStringNamespace:
             literal=literal,
         )
 
-    def head(self, n: int = 5) -> PandasExpr:
+    def slice(self, offset: int, length: int | None = None) -> PandasExpr:
         return reuse_series_namespace_implementation(
-            self._expr,
-            "str",
-            "head",
-            n,
+            self._expr, "str", "slice", offset, length
         )
 
     def to_datetime(self, format: str | None = None) -> PandasExpr:  # noqa: A002
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 61d61fbc..390f6add 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -513,6 +513,11 @@ class PandasSeriesStringNamespace:
     def __init__(self, series: PandasSeries) -> None:
         self._series = series
 
+    def starts_with(self, prefix: str) -> PandasSeries:
+        return self._series._from_series(
+            self._series._series.str.startswith(prefix),
+        )
+
     def ends_with(self, suffix: str) -> PandasSeries:
         return self._series._from_series(
             self._series._series.str.endswith(suffix),
@@ -523,9 +528,10 @@ class PandasSeriesStringNamespace:
             self._series._series.str.contains(pat=pattern, regex=not literal)
         )
 
-    def head(self, n: int = 5) -> PandasSeries:
+    def slice(self, offset: int, length: int | None = None) -> PandasSeries:
+        stop = offset + length if length else None
         return self._series._from_series(
-            self._series._series.str[:n],
+            self._series._series.str.slice(start=offset, stop=stop),
         )
 
     def to_datetime(self, format: str | None = None) -> PandasSeries:  # noqa: A002
diff --git a/narwhals/expression.py b/narwhals/expression.py
index d6829c3c..72e2c1ff 100644
--- a/narwhals/expression.py
+++ b/narwhals/expression.py
@@ -1532,22 +1532,106 @@ class ExprStringNamespace:
     def __init__(self, expr: Expr) -> None:
         self._expr = expr
 
+    def starts_with(self, prefix: str) -> Expr:
+        r"""
+        Check if string values start with a substring.
+
+        Arguments:
+            prefix: prefix substring
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"fruits": ["apple", "mango", None]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.with_columns(has_prefix=nw.col("fruits").str.starts_with("app"))
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+              fruits has_prefix
+            0  apple       True
+            1  mango      False
+            2   None       None
+
+            >>> func(df_pl)
+            shape: (3, 2)
+            ┌────────┬────────────┐
+            │ fruits ┆ has_prefix │
+            │ ---    ┆ ---        │
+            │ str    ┆ bool       │
+            ╞════════╪════════════╡
+            │ apple  ┆ true       │
+            │ mango  ┆ false      │
+            │ null   ┆ null       │
+            └────────┴────────────┘
+        """
+        return self._expr.__class__(
+            lambda plx: self._expr._call(plx).str.starts_with(prefix)
+        )
+
     def ends_with(self, suffix: str) -> Expr:
+        r"""
+        Check if string values end with a substring.
+
+        Arguments:
+            suffix: suffix substring
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"fruits": ["apple", "mango", None]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.with_columns(has_suffix=nw.col("fruits").str.ends_with("ngo"))
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+              fruits has_suffix
+            0  apple      False
+            1  mango       True
+            2   None       None
+
+            >>> func(df_pl)
+            shape: (3, 2)
+            ┌────────┬────────────┐
+            │ fruits ┆ has_suffix │
+            │ ---    ┆ ---        │
+            │ str    ┆ bool       │
+            ╞════════╪════════════╡
+            │ apple  ┆ false      │
+            │ mango  ┆ true       │
+            │ null   ┆ null       │
+            └────────┴────────────┘
+        """
         return self._expr.__class__(
             lambda plx: self._expr._call(plx).str.ends_with(suffix)
         )
 
     def contains(self, pattern: str, *, literal: bool = False) -> Expr:
-        """
+        r"""
         Check if string contains a substring that matches a pattern.
 
         Arguments:
             pattern: A Character sequence or valid regular expression pattern.
-
             literal: If True, treats the pattern as a literal string.
                      If False, assumes the pattern is a regular expression.
 
-        Example:
+        Examples:
             >>> import pandas as pd
             >>> import polars as pl
             >>> import narwhals as nw
@@ -1589,19 +1673,96 @@ class ExprStringNamespace:
             │ dove              ┆ false         ┆ true                   ┆ false         │
             │ null              ┆ null          ┆ null                   ┆ null          │
             └───────────────────┴───────────────┴────────────────────────┴───────────────┘
-
         """
 
         return self._expr.__class__(
             lambda plx: self._expr._call(plx).str.contains(pattern, literal=literal)
         )
 
-    def head(self, n: int = 5) -> Expr:
+    def slice(self, offset: int, length: int | None = None) -> Expr:
+        r"""
+        Create subslices of the string values of an expression.
+
+        Arguments:
+            offset: Start index. Negative indexing is supported.
+            length: Length of the slice. If set to `None` (default), the slice is taken to the
+                end of the string.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"s": ["pear", None, "papaya", "dragonfruit"]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.with_columns(s_sliced=nw.col("s").str.slice(4, length=3))
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)  # doctest: +NORMALIZE_WHITESPACE
+                         s s_sliced
+            0         pear
+            1         None     None
+            2       papaya       ya
+            3  dragonfruit      onf
+
+            >>> func(df_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (4, 2)
+            ┌─────────────┬──────────┐
+            │ s           ┆ s_sliced │
+            │ ---         ┆ ---      │
+            │ str         ┆ str      │
+            ╞═════════════╪══════════╡
+            │ pear        ┆          │
+            │ null        ┆ null     │
+            │ papaya      ┆ ya       │
+            │ dragonfruit ┆ onf      │
+            └─────────────┴──────────┘
+
+            Using negative indexes:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.with_columns(s_sliced=nw.col("s").str.slice(-3))
+
+            >>> func(df_pd)  # doctest: +NORMALIZE_WHITESPACE
+                         s s_sliced
+            0         pear      ear
+            1         None     None
+            2       papaya      aya
+            3  dragonfruit      uit
+
+            >>> func(df_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (4, 2)
+            ┌─────────────┬──────────┐
+            │ s           ┆ s_sliced │
+            │ ---         ┆ ---      │
+            │ str         ┆ str      │
+            ╞═════════════╪══════════╡
+            │ pear        ┆ ear      │
+            │ null        ┆ null     │
+            │ papaya      ┆ aya      │
+            │ dragonfruit ┆ uit      │
+            └─────────────┴──────────┘
         """
+        return self._expr.__class__(
+            lambda plx: self._expr._call(plx).str.slice(offset=offset, length=length)
+        )
+
+    def head(self, n: int = 5) -> Expr:
+        r"""
         Take the first n elements of each string.
 
         Arguments:
-            n: Number of elements to take.
+            n: Number of elements to take. Negative indexing is **not** supported.
+
+        Notes:
+            If the length of the string has fewer than `n` characters, the full string is returned.
 
         Examples:
             >>> import pandas as pd
@@ -1625,6 +1786,7 @@ class ExprStringNamespace:
             1      taata       taata
             2  taatatata       taata
             3    zukkyun       zukky
+
             >>> func(df_pl)
             shape: (4, 2)
             ┌───────────┬─────────────┐
@@ -1638,13 +1800,55 @@ class ExprStringNamespace:
             │ zukkyun   ┆ zukky       │
             └───────────┴─────────────┘
         """
+        return self._expr.__class__(lambda plx: self._expr._call(plx).str.slice(0, n))
+
+    def tail(self, n: int = 5) -> Expr:
+        r"""
+        Take the last n elements of each string.
+
+        Arguments:
+            n: Number of elements to take. Negative indexing is **not** supported.
+
+        Notes:
+            If the length of the string has fewer than `n` characters, the full string is returned.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"lyrics": ["Atatata", "taata", "taatatata", "zukkyun"]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
 
-        def func(plx: Any) -> Any:
-            if plx is get_polars():
-                return self._expr._call(plx).str.slice(0, n)
-            return self._expr._call(plx).str.head(n)
+            We define a dataframe-agnostic function:
 
-        return self._expr.__class__(func)
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.with_columns(lyrics_tail=nw.col("lyrics").str.tail())
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+                  lyrics lyrics_tail
+            0    Atatata       atata
+            1      taata       taata
+            2  taatatata       atata
+            3    zukkyun       kkyun
+
+            >>> func(df_pl)
+            shape: (4, 2)
+            ┌───────────┬─────────────┐
+            │ lyrics    ┆ lyrics_tail │
+            │ ---       ┆ ---         │
+            │ str       ┆ str         │
+            ╞═══════════╪═════════════╡
+            │ Atatata   ┆ atata       │
+            │ taata     ┆ taata       │
+            │ taatatata ┆ atata       │
+            │ zukkyun   ┆ kkyun       │
+            └───────────┴─────────────┘
+        """
+        return self._expr.__class__(lambda plx: self._expr._call(plx).str.slice(-n))
 
     def to_datetime(self, format: str) -> Expr:  # noqa: A002
         """
diff --git a/narwhals/series.py b/narwhals/series.py
index 1adc55e4..2a80d8cc 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -1720,20 +1720,96 @@ class SeriesStringNamespace:
     def __init__(self, series: Series) -> None:
         self._series = series
 
+    def starts_with(self, prefix: str) -> Series:
+        r"""
+        Check if string values start with a substring.
+
+        Arguments:
+            prefix: prefix substring
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = ["apple", "mango", None]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify(allow_series=True)
+            ... def func(series):
+            ...     return series.str.starts_with("app")
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0     True
+            1    False
+            2     None
+            dtype: object
+
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3,)
+            Series: '' [bool]
+            [
+               true
+               false
+               null
+            ]
+        """
+        return self._series.__class__(self._series._series.str.starts_with(prefix))
+
     def ends_with(self, suffix: str) -> Series:
+        r"""
+        Check if string values end with a substring.
+
+        Arguments:
+            suffix: suffix substring
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = ["apple", "mango", None]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify(allow_series=True)
+            ... def func(series):
+            ...     return series.str.ends_with("ngo")
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    False
+            1     True
+            2     None
+            dtype: object
+
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (3,)
+            Series: '' [bool]
+            [
+               false
+               true
+               null
+            ]
+        """
         return self._series.__class__(self._series._series.str.ends_with(suffix))
 
     def contains(self, pattern: str, *, literal: bool = False) -> Series:
-        """
+        r"""
         Check if string contains a substring that matches a pattern.
 
         Arguments:
             pattern: A Character sequence or valid regular expression pattern.
-
             literal: If True, treats the pattern as a literal string.
                      If False, assumes the pattern is a regular expression.
 
-        Example:
+        Examples:
             >>> import pandas as pd
             >>> import polars as pl
             >>> import narwhals as nw
@@ -1756,6 +1832,7 @@ class SeriesStringNamespace:
             3     True
             4     None
             dtype: object
+
             >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
             shape: (5,)
             Series: '' [bool]
@@ -1771,12 +1848,86 @@ class SeriesStringNamespace:
             self._series._series.str.contains(pattern, literal=literal)
         )
 
-    def head(self, n: int = 5) -> Series:
+    def slice(self, offset: int, length: int | None = None) -> Series:
+        r"""
+        Create subslices of the string values of a Series.
+
+        Arguments:
+            offset: Start index. Negative indexing is supported.
+            length: Length of the slice. If set to `None` (default), the slice is taken to the
+                end of the string.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = ["pear", None, "papaya", "dragonfruit"]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify(series_only=True)
+            ... def func(s):
+            ...     return s.str.slice(4, length=3)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)  # doctest: +NORMALIZE_WHITESPACE
+            0
+            1    None
+            2      ya
+            3     onf
+            dtype: object
+
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (4,)
+            Series: '' [str]
+            [
+               ""
+               null
+               "ya"
+               "onf"
+            ]
+
+            Using negative indexes:
+
+            >>> @nw.narwhalify(series_only=True)
+            ... def func(s):
+            ...     return s.str.slice(-3)
+
+            >>> func(s_pd)  # doctest: +NORMALIZE_WHITESPACE
+            0     ear
+            1    None
+            2     aya
+            3     uit
+            dtype: object
+
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (4,)
+            Series: '' [str]
+            [
+                "ear"
+                null
+                "aya"
+                "uit"
+            ]
         """
+        return self._series.__class__(
+            self._series._series.str.slice(offset=offset, length=length)
+        )
+
+    def head(self, n: int = 5) -> Series:
+        r"""
         Take the first n elements of each string.
 
         Arguments:
-            n: Number of elements to take.
+            n: Number of elements to take. Negative indexing is supported (see note (1.))
+
+        Notes:
+            1. When the `n` input is negative, `head` returns characters up to the n-th from the end of the string.
+                For example, if `n = -3`, then all characters except the last three are returned.
+            2. If the length of the string has fewer than `n` characters, the full string is returned.
 
         Examples:
             >>> import pandas as pd
@@ -1810,9 +1961,53 @@ class SeriesStringNamespace:
                "zukky"
             ]
         """
-        if self._series._is_polars:
-            return self._series.__class__(self._series._series.str.slice(0, n))
-        return self._series.__class__(self._series._series.str.head(n))
+        return self._series.__class__(self._series._series.str.slice(0, n))
+
+    def tail(self, n: int = 5) -> Series:
+        r"""
+        Take the last n elements of each string.
+
+        Arguments:
+            n: Number of elements to take. Negative indexing is supported (see note (1.))
+
+        Notes:
+            1. When the `n` input is negative, `tail` returns characters starting from the n-th from the beginning of
+                the string. For example, if `n = -3`, then all characters except the first three are returned.
+            2. If the length of the string has fewer than `n` characters, the full string is returned.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> lyrics = ["Atatata", "taata", "taatatata", "zukkyun"]
+            >>> s_pd = pd.Series(lyrics)
+            >>> s_pl = pl.Series(lyrics)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify(series_only=True)
+            ... def func(s):
+            ...     return s.str.tail()
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            0    atata
+            1    taata
+            2    atata
+            3    kkyun
+            dtype: object
+            >>> func(s_pl)  # doctest: +NORMALIZE_WHITESPACE
+            shape: (4,)
+            Series: '' [str]
+            [
+               "atata"
+               "taata"
+               "atata"
+               "kkyun"
+            ]
+        """
+        return self._series.__class__(self._series._series.str.slice(-n))
 
 
 class SeriesDateTimeNamespace:


Test patch:
diff --git a/tests/expr/str/slice_test.py b/tests/expr/str/slice_test.py
new file mode 100644
index 00000000..a97cfc81
--- /dev/null
+++ b/tests/expr/str/slice_test.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {"a": ["fdas", "edfas"]}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize(
+    ("offset", "length", "expected"),
+    [(1, 2, {"a": ["da", "df"]}), (-2, None, {"a": ["as", "as"]})],
+)
+def test_str_slice(
+    constructor: Any, offset: int, length: int | None, expected: Any
+) -> None:
+    df = nw.from_native(constructor(data), eager_only=True)
+    result_frame = df.select(nw.col("a").str.slice(offset, length))
+    compare_dicts(result_frame, expected)
+
+    result_series = df["a"].str.slice(offset, length)
+    assert result_series.to_numpy().tolist() == expected["a"]
diff --git a/tests/expr/str/tail_test.py b/tests/expr/str/tail_test.py
new file mode 100644
index 00000000..1ab9d59e
--- /dev/null
+++ b/tests/expr/str/tail_test.py
@@ -0,0 +1,26 @@
+from typing import Any
+
+import pandas as pd
+import polars as pl
+import pytest
+
+import narwhals as nw
+from tests.utils import compare_dicts
+
+data = {
+    "a": ["foo", "bars"],
+}
+
+
+@pytest.mark.parametrize("constructor", [pd.DataFrame, pl.DataFrame])
+def test_str_tail(constructor: Any) -> None:
+    df = nw.from_native(constructor(data), eager_only=True)
+    expected = {
+        "a": ["foo", "ars"],
+    }
+
+    result_frame = df.select(nw.col("a").str.tail(3))
+    compare_dicts(result_frame, expected)
+
+    result_series = df["a"].str.tail(3)
+    assert result_series.to_numpy().tolist() == expected["a"]
diff --git a/tests/test_str.py b/tests/test_str.py
index 9e3da9f7..7e1ae49b 100644
--- a/tests/test_str.py
+++ b/tests/test_str.py
@@ -45,3 +45,23 @@ def test_ends_with(df_raw: Any) -> None:
         "a": [True, False],
     }
     compare_dicts(result_native, expected)
+
+
+@pytest.mark.parametrize(
+    "df_raw",
+    [df_pandas, df_polars, df_mpd],
+)
+def test_starts_with(df_raw: Any) -> None:
+    df = nw.LazyFrame(df_raw)
+    result = df.select(nw.col("a").str.starts_with("fda"))
+    result_native = nw.to_native(result)
+    expected = {
+        "a": [True, False],
+    }
+    compare_dicts(result_native, expected)
+    result = df.select(df.collect()["a"].str.starts_with("fda"))
+    result_native = nw.to_native(result)
+    expected = {
+        "a": [True, False],
+    }
+    compare_dicts(result_native, expected)