Skip to content

Commit e3133de

Browse files
dangotbannedFBruzzesi
dangotbanned
and
FBruzzesi
authored
chore: Support @requires.backend_version in namespaces (#3127)
Co-authored-by: FBruzzesi <[email protected]> Co-authored-by: Francesco Bruzzesi <[email protected]>
1 parent ba40f30 commit e3133de

File tree

12 files changed

+139
-65
lines changed

12 files changed

+139
-65
lines changed

‎narwhals/_arrow/series_cat.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,14 @@
55
import pyarrow as pa
66

77
from narwhals._arrow.utils import ArrowSeriesNamespace
8+
from narwhals._compliant.any_namespace import CatNamespace
89

910
if TYPE_CHECKING:
1011
from narwhals._arrow.series import ArrowSeries
1112
from narwhals._arrow.typing import Incomplete
1213

1314

14-
class ArrowSeriesCatNamespace(ArrowSeriesNamespace):
15+
class ArrowSeriesCatNamespace(ArrowSeriesNamespace, CatNamespace["ArrowSeries"]):
1516
def get_categories(self) -> ArrowSeries:
1617
# NOTE: Should be `list[pa.DictionaryArray]`, but `DictionaryArray` has no attributes
1718
chunks: Incomplete = self.native.chunks

‎narwhals/_arrow/series_dt.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import pyarrow.compute as pc
77

88
from narwhals._arrow.utils import UNITS_DICT, ArrowSeriesNamespace, floordiv_compat, lit
9+
from narwhals._compliant.any_namespace import DateTimeNamespace
910
from narwhals._constants import (
1011
MS_PER_MINUTE,
1112
MS_PER_SECOND,
@@ -36,7 +37,9 @@
3637
IntoRhs: TypeAlias = int
3738

3839

39-
class ArrowSeriesDateTimeNamespace(ArrowSeriesNamespace):
40+
class ArrowSeriesDateTimeNamespace(
41+
ArrowSeriesNamespace, DateTimeNamespace["ArrowSeries"]
42+
):
4043
_TIMESTAMP_DATE_FACTOR: ClassVar[Mapping[TimeUnit, int]] = {
4144
"ns": NS_PER_SECOND,
4245
"us": US_PER_SECOND,

‎narwhals/_arrow/series_list.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,19 +6,19 @@
66
import pyarrow.compute as pc
77

88
from narwhals._arrow.utils import ArrowSeriesNamespace
9+
from narwhals._compliant.any_namespace import ListNamespace
910
from narwhals._utils import not_implemented
1011

1112
if TYPE_CHECKING:
1213
from narwhals._arrow.series import ArrowSeries
1314

1415

15-
class ArrowSeriesListNamespace(ArrowSeriesNamespace):
16+
class ArrowSeriesListNamespace(ArrowSeriesNamespace, ListNamespace["ArrowSeries"]):
1617
def len(self) -> ArrowSeries:
1718
return self.with_native(pc.list_value_length(self.native).cast(pa.uint32()))
1819

19-
unique = not_implemented()
20-
21-
contains = not_implemented()
22-
2320
def get(self, index: int) -> ArrowSeries:
2421
return self.with_native(pc.list_element(self.native, index))
22+
23+
unique = not_implemented()
24+
contains = not_implemented()

‎narwhals/_arrow/series_str.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,14 @@
77
import pyarrow.compute as pc
88

99
from narwhals._arrow.utils import ArrowSeriesNamespace, lit, parse_datetime_format
10+
from narwhals._compliant.any_namespace import StringNamespace
1011

1112
if TYPE_CHECKING:
1213
from narwhals._arrow.series import ArrowSeries
1314
from narwhals._arrow.typing import Incomplete
1415

1516

16-
class ArrowSeriesStringNamespace(ArrowSeriesNamespace):
17+
class ArrowSeriesStringNamespace(ArrowSeriesNamespace, StringNamespace["ArrowSeries"]):
1718
def len_chars(self) -> ArrowSeries:
1819
return self.with_native(pc.utf8_length(self.native))
1920

‎narwhals/_arrow/series_struct.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,12 @@
55
import pyarrow.compute as pc
66

77
from narwhals._arrow.utils import ArrowSeriesNamespace
8+
from narwhals._compliant.any_namespace import StructNamespace
89

910
if TYPE_CHECKING:
1011
from narwhals._arrow.series import ArrowSeries
1112

1213

13-
class ArrowSeriesStructNamespace(ArrowSeriesNamespace):
14+
class ArrowSeriesStructNamespace(ArrowSeriesNamespace, StructNamespace["ArrowSeries"]):
1415
def field(self, name: str) -> ArrowSeries:
1516
return self.with_native(pc.struct_field(self.native, name)).alias(name)

‎narwhals/_compliant/any_namespace.py‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,30 +2,40 @@
22

33
from __future__ import annotations
44

5-
from typing import TYPE_CHECKING, Protocol
5+
from typing import TYPE_CHECKING, ClassVar, Protocol
66

77
from narwhals._utils import CompliantT_co, _StoresCompliant
88

99
if TYPE_CHECKING:
1010
from typing import Callable
1111

12+
from narwhals._compliant.typing import Accessor
1213
from narwhals.typing import NonNestedLiteral, TimeUnit
1314

1415
__all__ = [
1516
"CatNamespace",
1617
"DateTimeNamespace",
1718
"ListNamespace",
1819
"NameNamespace",
20+
"NamespaceAccessor",
1921
"StringNamespace",
2022
"StructNamespace",
2123
]
2224

2325

24-
class CatNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
26+
class NamespaceAccessor(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
27+
_accessor: ClassVar[Accessor]
28+
29+
30+
class CatNamespace(NamespaceAccessor[CompliantT_co], Protocol[CompliantT_co]):
31+
_accessor: ClassVar[Accessor] = "cat"
32+
2533
def get_categories(self) -> CompliantT_co: ...
2634

2735

2836
class DateTimeNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
37+
_accessor: ClassVar[Accessor] = "dt"
38+
2939
def to_string(self, format: str) -> CompliantT_co: ...
3040
def replace_time_zone(self, time_zone: str | None) -> CompliantT_co: ...
3141
def convert_time_zone(self, time_zone: str) -> CompliantT_co: ...
@@ -52,15 +62,17 @@ def offset_by(self, by: str) -> CompliantT_co: ...
5262

5363

5464
class ListNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
55-
def get(self, index: int) -> CompliantT_co: ...
65+
_accessor: ClassVar[Accessor] = "list"
5666

67+
def get(self, index: int) -> CompliantT_co: ...
5768
def len(self) -> CompliantT_co: ...
58-
5969
def unique(self) -> CompliantT_co: ...
6070
def contains(self, item: NonNestedLiteral) -> CompliantT_co: ...
6171

6272

6373
class NameNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
74+
_accessor: ClassVar[Accessor] = "name"
75+
6476
def keep(self) -> CompliantT_co: ...
6577
def map(self, function: Callable[[str], str]) -> CompliantT_co: ...
6678
def prefix(self, prefix: str) -> CompliantT_co: ...
@@ -70,6 +82,8 @@ def to_uppercase(self) -> CompliantT_co: ...
7082

7183

7284
class StringNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
85+
_accessor: ClassVar[Accessor] = "str"
86+
7387
def len_chars(self) -> CompliantT_co: ...
7488
def replace(
7589
self, pattern: str, value: str, *, literal: bool, n: int
@@ -91,4 +105,6 @@ def zfill(self, width: int) -> CompliantT_co: ...
91105

92106

93107
class StructNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):
108+
_accessor: ClassVar[Accessor] = "struct"
109+
94110
def field(self, name: str) -> CompliantT_co: ...

‎narwhals/_compliant/typing.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,3 +204,8 @@ class ScalarKwargs(TypedDict, total=False):
204204
- https://github.com/narwhals-dev/narwhals/issues/2526
205205
- https://github.com/narwhals-dev/narwhals/issues/2660
206206
"""
207+
208+
Accessor: TypeAlias = Literal[
209+
"arr", "cat", "dt", "list", "meta", "name", "str", "bin", "struct"
210+
]
211+
"""`{Expr,Series}` method namespace accessor name."""

‎narwhals/_polars/expr.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import annotations
22

3-
from typing import TYPE_CHECKING, Any, Callable, Literal
3+
from typing import TYPE_CHECKING, Any, Callable, ClassVar, Literal
44

55
import polars as pl
66

@@ -22,6 +22,7 @@
2222

2323
from typing_extensions import Self
2424

25+
from narwhals._compliant.typing import Accessor
2526
from narwhals._expression_parsing import ExprKind, ExprMetadata
2627
from narwhals._polars.dataframe import Method
2728
from narwhals._polars.namespace import PolarsNamespace
@@ -400,17 +401,11 @@ class PolarsExprDateTimeNamespace(
400401
class PolarsExprStringNamespace(
401402
PolarsExprNamespace, PolarsStringNamespace[PolarsExpr, pl.Expr]
402403
):
404+
@requires.backend_version((0, 20, 5))
403405
def zfill(self, width: int) -> PolarsExpr:
404406
backend_version = self.compliant._backend_version
405407
native_result = self.native.str.zfill(width)
406408

407-
if backend_version < (0, 20, 5): # pragma: no cover
408-
# Reason:
409-
# `TypeError: argument 'length': 'Expr' object cannot be interpreted as an integer`
410-
# in `native_expr.str.slice(1, length)`
411-
msg = "`zfill` is only available in 'polars>=0.20.5', found version '0.20.4'."
412-
raise NotImplementedError(msg)
413-
414409
if backend_version <= (1, 30, 0):
415410
length = self.native.str.len_chars()
416411
less_than_width = length < width
@@ -435,7 +430,7 @@ class PolarsExprCatNamespace(
435430

436431

437432
class PolarsExprNameNamespace(PolarsExprNamespace):
438-
_accessor = "name"
433+
_accessor: ClassVar[Accessor] = "name"
439434
keep: Method[PolarsExpr]
440435
map: Method[PolarsExpr]
441436
prefix: Method[PolarsExpr]

‎narwhals/_polars/typing.py‎

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,10 @@
11
from __future__ import annotations # pragma: no cover
22

3-
from typing import (
4-
TYPE_CHECKING, # pragma: no cover
5-
Union, # pragma: no cover
6-
)
3+
from typing import TYPE_CHECKING # pragma: no cover
74

85
if TYPE_CHECKING:
9-
import sys
10-
from typing import Literal, TypeVar
11-
12-
if sys.version_info >= (3, 10):
13-
from typing import TypeAlias
14-
else:
15-
from typing_extensions import TypeAlias
6+
from typing import TypeVar
167

178
from narwhals._polars.dataframe import PolarsDataFrame, PolarsLazyFrame
18-
from narwhals._polars.expr import PolarsExpr
19-
from narwhals._polars.series import PolarsSeries
209

21-
IntoPolarsExpr: TypeAlias = Union[PolarsExpr, PolarsSeries]
2210
FrameT = TypeVar("FrameT", PolarsDataFrame, PolarsLazyFrame)
23-
NativeAccessor: TypeAlias = Literal[
24-
"arr", "cat", "dt", "list", "meta", "name", "str", "bin", "struct"
25-
]

‎narwhals/_polars/utils.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,10 @@
3030

3131
from typing_extensions import TypeIs
3232

33+
from narwhals._compliant.typing import Accessor
3334
from narwhals._polars.dataframe import Method
3435
from narwhals._polars.expr import PolarsExpr
3536
from narwhals._polars.series import PolarsSeries
36-
from narwhals._polars.typing import NativeAccessor
3737
from narwhals.dtypes import DType
3838
from narwhals.typing import IntoDType
3939

@@ -261,7 +261,7 @@ class PolarsAnyNamespace(
261261
_StoresNative[NativeT_co],
262262
Protocol[CompliantT_co, NativeT_co],
263263
):
264-
_accessor: ClassVar[NativeAccessor]
264+
_accessor: ClassVar[Accessor]
265265

266266
def __getattr__(self, attr: str) -> Callable[..., CompliantT_co]:
267267
def func(*args: Any, **kwargs: Any) -> CompliantT_co:
@@ -273,7 +273,7 @@ def func(*args: Any, **kwargs: Any) -> CompliantT_co:
273273

274274

275275
class PolarsDateTimeNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
276-
_accessor: ClassVar[NativeAccessor] = "dt"
276+
_accessor: ClassVar[Accessor] = "dt"
277277

278278
def truncate(self, every: str) -> CompliantT:
279279
# Ensure consistent error message is raised.
@@ -309,7 +309,7 @@ def offset_by(self, by: str) -> CompliantT:
309309

310310

311311
class PolarsStringNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
312-
_accessor: ClassVar[NativeAccessor] = "str"
312+
_accessor: ClassVar[Accessor] = "str"
313313

314314
# NOTE: Use `abstractmethod` if we have defs to implement, but also `Method` usage
315315
@abc.abstractmethod
@@ -331,12 +331,12 @@ def zfill(self, width: int) -> CompliantT: ...
331331

332332

333333
class PolarsCatNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
334-
_accessor: ClassVar[NativeAccessor] = "cat"
334+
_accessor: ClassVar[Accessor] = "cat"
335335
get_categories: Method[CompliantT]
336336

337337

338338
class PolarsListNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
339-
_accessor: ClassVar[NativeAccessor] = "list"
339+
_accessor: ClassVar[Accessor] = "list"
340340

341341
@abc.abstractmethod
342342
def len(self) -> CompliantT: ...
@@ -347,5 +347,5 @@ def len(self) -> CompliantT: ...
347347

348348

349349
class PolarsStructNamespace(PolarsAnyNamespace[CompliantT, NativeT_co]):
350-
_accessor: ClassVar[NativeAccessor] = "struct"
350+
_accessor: ClassVar[Accessor] = "struct"
351351
field: Method[CompliantT]

0 commit comments

Comments
 (0)