Skip to content

Commit

Permalink
feat[python]: api consistency; allow fluent calling pattern everywhere (
Browse files Browse the repository at this point in the history 
#4761)
  • Loading branch information
@alexander-beedie
alexander-beedie committed Sep 7, 2022
1 parent 56d2e1d commit ca7cda5
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 154 deletions.
85 changes: 22 additions & 63 deletions py-polars/polars/internals/dataframe/frame.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2295,7 +2295,7 @@ def reverse(self: DF) -> DF:
"""
return self.select(pli.col("*").reverse())

def rename(self, mapping: dict[str, str]) -> DataFrame:
def rename(self: DF, mapping: dict[str, str]) -> DF | DataFrame:
"""
Rename column names.
Expand Down Expand Up @@ -2326,7 +2326,7 @@ def rename(self, mapping: dict[str, str]) -> DataFrame:
"""
return self.lazy().rename(mapping).collect(no_optimization=True)

def insert_at_idx(self, index: int, series: pli.Series) -> None:
def insert_at_idx(self: DF, index: int, series: pli.Series) -> DF:
"""
Insert a Series at a certain column index. This operation is in place.
Expand All @@ -2341,8 +2341,7 @@ def insert_at_idx(self, index: int, series: pli.Series) -> None:
--------
>>> df = pl.DataFrame({"foo": [1, 2, 3], "bar": [4, 5, 6]})
>>> s = pl.Series("baz", [97, 98, 99])
>>> df.insert_at_idx(1, s) # returns None
>>> df
>>> df.insert_at_idx(1, s)
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ baz ┆ bar │
Expand All @@ -2365,7 +2364,6 @@ def insert_at_idx(self, index: int, series: pli.Series) -> None:
... )
>>> s = pl.Series("d", [-2.5, 15, 20.5, 0])
>>> df.insert_at_idx(3, s)
>>> df
shape: (4, 4)
┌─────┬──────┬───────┬──────┐
│ a ┆ b ┆ c ┆ d │
Expand All @@ -2385,6 +2383,7 @@ def insert_at_idx(self, index: int, series: pli.Series) -> None:
if index < 0:
index = len(self.columns) + index
self._df.insert_at_idx(index, series._s)
return self

def filter(
self,
Expand Down Expand Up @@ -2525,7 +2524,7 @@ def find_idx_by_name(self, name: str) -> int:
"""
return self._df.find_idx_by_name(name)

def replace_at_idx(self, index: int, series: pli.Series) -> None:
def replace_at_idx(self: DF, index: int, series: pli.Series) -> DF:
"""
Replace a column at an index location.
Expand All @@ -2545,9 +2544,8 @@ def replace_at_idx(self, index: int, series: pli.Series) -> None:
... "ham": ["a", "b", "c"],
... }
... )
>>> x = pl.Series("apple", [10, 20, 30])
>>> df.replace_at_idx(0, x)
>>> df
>>> s = pl.Series("apple", [10, 20, 30])
>>> df.replace_at_idx(0, s)
shape: (3, 3)
┌───────┬─────┬─────┐
│ apple ┆ bar ┆ ham │
Expand All @@ -2565,13 +2563,14 @@ def replace_at_idx(self, index: int, series: pli.Series) -> None:
if index < 0:
index = len(self.columns) + index
self._df.replace_at_idx(index, series._s)
return self

def sort(
self,
self: DF,
by: str | pli.Expr | list[str] | list[pli.Expr],
reverse: bool | list[bool] = False,
nulls_last: bool = False,
) -> DataFrame:
) -> DF | DataFrame:
"""
Sort the DataFrame by column.
Expand Down Expand Up @@ -2672,7 +2671,7 @@ def frame_equal(self, other: DataFrame, null_equal: bool = True) -> bool:
"""
return self._df.frame_equal(other._df, null_equal)

def replace(self, column: str, new_col: pli.Series) -> None:
def replace(self: DF, column: str, new_col: pli.Series) -> DF:
"""
Replace a column by a new Series.
Expand All @@ -2688,7 +2687,6 @@ def replace(self, column: str, new_col: pli.Series) -> None:
>>> df = pl.DataFrame({"foo": [1, 2, 3], "bar": [4, 5, 6]})
>>> s = pl.Series([10, 20, 30])
>>> df.replace("foo", s) # works in-place!
>>> df
shape: (3, 2)
┌─────┬─────┐
│ foo ┆ bar │
Expand All @@ -2704,6 +2702,7 @@ def replace(self, column: str, new_col: pli.Series) -> None:
"""
self._df.replace(column, new_col._s)
return self

def slice(self: DF, offset: int, length: int | None = None) -> DF:
"""
Expand Down Expand Up @@ -4024,27 +4023,11 @@ def with_column(self, column: pli.Series | pli.Expr) -> DataFrame:
else:
return self._from_pydf(self._df.with_column(column._s))

@overload
def hstack(
self: DF,
columns: list[pli.Series] | DataFrame,
in_place: Literal[False] = False,
) -> DF:
...

@overload
def hstack(
self: DF,
columns: list[pli.Series] | DataFrame,
in_place: Literal[True],
) -> None:
...

def hstack(
self: DF,
columns: list[pli.Series] | DataFrame,
in_place: bool = False,
) -> DF | None:
) -> DF:
"""
Return a new DataFrame grown horizontally by stacking multiple Series to it.
Expand Down Expand Up @@ -4084,23 +4067,11 @@ def hstack(
columns = columns.get_columns()
if in_place:
self._df.hstack_mut([s._s for s in columns])
return None
return self
else:
return self._from_pydf(self._df.hstack([s._s for s in columns]))

@overload
def vstack(self, df: DataFrame, in_place: Literal[True]) -> None:
...

@overload
def vstack(self: DF, df: DataFrame, in_place: Literal[False] = ...) -> DF:
...

@overload
def vstack(self: DF, df: DataFrame, in_place: bool) -> DF | None:
...

def vstack(self: DF, df: DataFrame, in_place: bool = False) -> DF | None:
def vstack(self: DF, df: DataFrame, in_place: bool = False) -> DF:
"""
Grow this DataFrame vertically by stacking a DataFrame to it.
Expand Down Expand Up @@ -4146,16 +4117,16 @@ def vstack(self: DF, df: DataFrame, in_place: bool = False) -> DF | None:
"""
if in_place:
self._df.vstack_mut(df._df)
return None
return self
else:
return self._from_pydf(self._df.vstack(df._df))

def extend(self, other: DataFrame) -> None:
def extend(self: DF, other: DF) -> DF:
"""
Extend the memory backed by this `DataFrame` with the values from `other`.
Different from `vstack` which adds the chunks from `other` to the chunks of this
`DataFrame` `extent` appends the data from `other` to the underlying memory
`DataFrame` `extend` appends the data from `other` to the underlying memory
locations and thus may cause a reallocation.
If this does not cause a reallocation, the resulting data structure will not
Expand All @@ -4178,8 +4149,7 @@ def extend(self, other: DataFrame) -> None:
--------
>>> df1 = pl.DataFrame({"foo": [1, 2, 3], "bar": [4, 5, 6]})
>>> df2 = pl.DataFrame({"foo": [10, 20, 30], "bar": [40, 50, 60]})
>>> df1.extend(df2) # returns None
>>> df1
>>> df1.extend(df2)
shape: (6, 2)
┌─────┬─────┐
│ foo ┆ bar │
Expand All @@ -4201,6 +4171,7 @@ def extend(self, other: DataFrame) -> None:
"""
self._df.extend(other._df)
return self

def drop(self: DF, name: str | list[str]) -> DF:
"""
Expand Down Expand Up @@ -5943,19 +5914,7 @@ def rows(self) -> list[tuple[object, ...]]:
"""
return self._df.row_tuples()

@overload
def shrink_to_fit(self: DF, in_place: Literal[False] = ...) -> DF:
...

@overload
def shrink_to_fit(self, in_place: Literal[True]) -> None:
...

@overload
def shrink_to_fit(self: DF, in_place: bool) -> DF | None:
...

def shrink_to_fit(self: DF, in_place: bool = False) -> DF | None:
def shrink_to_fit(self: DF, in_place: bool = False) -> DF:
"""
Shrink DataFrame memory usage.
Expand All @@ -5964,7 +5923,7 @@ def shrink_to_fit(self: DF, in_place: bool = False) -> DF | None:
"""
if in_place:
self._df.shrink_to_fit()
return None
return self
else:
df = self.clone()
df._df.shrink_to_fit()
Expand Down

0 comments on commit ca7cda5

Please sign in to comment.