Skip to content

Commit d4e07e8

Browse files
committed
refactor(DRAFT): Make PolarsDataFrame compliant
- Technically works, but I'd like to avoid the `Any` - Maybe adding a generic callable alias, that gets subscribed w/ `Self`?
1 parent c4095bc commit d4e07e8

File tree

1 file changed

+30
-3
lines changed

1 file changed

+30
-3
lines changed

narwhals/_polars/dataframe.py

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,34 @@
3939
T = TypeVar("T")
4040

4141

42+
# TODO @dangotbanned: Want to tell the type checker that `__getattr__` will satisfy specfic methods
43+
# - Could these have a narrower annotation?
4244
class PolarsDataFrame:
45+
clone: Any
46+
collect: Any
47+
drop_nulls: Any
48+
estimated_size: Any
49+
filter: Any
50+
gather_every: Any
51+
item: Any
52+
iter_rows: Any
53+
is_unique: Any
54+
join: Any
55+
join_asof: Any
56+
rename: Any
57+
row: Any
58+
rows: Any
59+
sample: Any
60+
select: Any
61+
sort: Any
62+
to_arrow: Any
63+
to_numpy: Any
64+
to_pandas: Any
65+
unique: Any
66+
with_columns: Any
67+
write_csv: Any
68+
write_parquet: Any
69+
4370
def __init__(
4471
self: Self,
4572
df: pl.DataFrame,
@@ -321,16 +348,16 @@ def with_row_index(self: Self, name: str) -> Self:
321348
return self._from_native_frame(self._native_frame.with_row_count(name))
322349
return self._from_native_frame(self._native_frame.with_row_index(name))
323350

324-
def drop(self: Self, columns: list[str], strict: bool) -> Self: # noqa: FBT001
351+
def drop(self: Self, columns: Sequence[str], strict: bool) -> Self: # noqa: FBT001
325352
to_drop = parse_columns_to_drop(
326353
compliant_frame=self, columns=columns, strict=strict
327354
)
328355
return self._from_native_frame(self._native_frame.drop(to_drop))
329356

330357
def unpivot(
331358
self: Self,
332-
on: list[str] | None,
333-
index: list[str] | None,
359+
on: Sequence[str] | None,
360+
index: Sequence[str] | None,
334361
variable_name: str,
335362
value_name: str,
336363
) -> Self:

0 commit comments

Comments
 (0)