Skip to content

Commit

Permalink
simplify
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Mar 10, 2024
1 parent df7e7f2 commit f8cb4a5
Show file tree
Hide file tree
Showing 10 changed files with 195 additions and 178 deletions.
20 changes: 10 additions & 10 deletions narwhals/pandas_like/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,12 @@
from typing_extensions import Self

from narwhals.pandas_like.group_by import PdxGroupBy
from narwhals.pandas_like.series import PandasLikeSeries
from narwhals.pandas_like.series import PandasSeries
from narwhals.spec import DType
from narwhals.spec import IntoExpr


class PdxDataFrame(DataFrameProtocol):
class PandasDataFrame(DataFrameProtocol):
# --- not in the spec ---
def __init__(
self,
Expand Down Expand Up @@ -81,14 +81,14 @@ def _from_dataframe(self, df: Any) -> Self:
is_lazy=self._is_lazy,
)

def __getitem__(self, column_name: str) -> PandasLikeSeries:
from narwhals.pandas_like.series import PandasLikeSeries
def __getitem__(self, column_name: str) -> PandasSeries:
from narwhals.pandas_like.series import PandasSeries

if not self._is_eager:
raise RuntimeError(
"DataFrame.__getitem__ can only be called when it was instantiated with `is_eager=True`"
)
return PandasLikeSeries(
return PandasSeries(
self._dataframe.loc[:, column_name],
implementation=self._implementation,
)
Expand Down Expand Up @@ -164,12 +164,12 @@ def sort(
)

# --- convert ---
def collect(self) -> PdxDataFrame:
def collect(self) -> PandasDataFrame:
if not self._is_lazy:
raise RuntimeError(
"DataFrame.collect can only be called when it was instantiated with `is_lazy=True`"
)
return PdxDataFrame(
return PandasDataFrame(
self._dataframe,
implementation=self._implementation,
is_eager=True,
Expand Down Expand Up @@ -247,15 +247,15 @@ def shape(self) -> tuple[int, int]:
)
return self._dataframe.shape # type: ignore[no-any-return]

def iter_columns(self) -> Iterable[PandasLikeSeries]:
from narwhals.pandas_like.series import PandasLikeSeries
def iter_columns(self) -> Iterable[PandasSeries]:
from narwhals.pandas_like.series import PandasSeries

if not self._is_eager:
raise RuntimeError(
"DataFrame.iter_columns can only be called when it was instantiated with `is_eager=True`"
)
return (
PandasLikeSeries(self._dataframe[col], implementation=self._implementation)
PandasSeries(self._dataframe[col], implementation=self._implementation)
for col in self.columns
)

Expand Down
12 changes: 6 additions & 6 deletions narwhals/pandas_like/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,22 +4,22 @@
from typing import Any
from typing import Callable

from narwhals.pandas_like.series import PandasLikeSeries
from narwhals.pandas_like.series import PandasSeries
from narwhals.pandas_like.utils import register_expression_call
from narwhals.spec import Expr as ExprProtocol
from narwhals.spec import ExprStringNamespace as ExprStringNamespaceProtocol

if TYPE_CHECKING:
from typing_extensions import Self

from narwhals.pandas_like.dataframe import PdxDataFrame
from narwhals.pandas_like.dataframe import PandasDataFrame
from narwhals.pandas_like.dtypes import DType


class Expr(ExprProtocol):
def __init__( # noqa: PLR0913
self,
call: Callable[[PdxDataFrame], list[PandasLikeSeries]],
call: Callable[[PandasDataFrame], list[PandasSeries]],
*,
depth: int,
function_name: str,
Expand Down Expand Up @@ -50,7 +50,7 @@ def from_column_names(
) -> Self:
return cls(
lambda df: [
PandasLikeSeries(
PandasSeries(
df._dataframe.loc[:, column_name],
implementation=implementation,
)
Expand Down Expand Up @@ -212,7 +212,7 @@ def ends_with(self, suffix: str) -> Expr:

return Expr(
lambda df: [
PandasLikeSeries(
PandasSeries(
series.series.str.endswith(suffix),
implementation=df._implementation,
)
Expand All @@ -228,7 +228,7 @@ def ends_with(self, suffix: str) -> Expr:
def strip_chars(self, characters: str = " ") -> Expr:
return Expr(
lambda df: [
PandasLikeSeries(
PandasSeries(
series.series.str.strip(characters),
implementation=df._implementation,
)
Expand Down
20 changes: 10 additions & 10 deletions narwhals/pandas_like/group_by.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,13 @@
from narwhals.utils import remove_prefix

if TYPE_CHECKING:
from narwhals.pandas_like.dataframe import PdxDataFrame
from narwhals.pandas_like.dataframe import PandasDataFrame
from narwhals.pandas_like.expr import Expr


class PdxGroupBy(GroupByProtocol):
def __init__(
self, df: PdxDataFrame, keys: list[str], *, is_eager: bool, is_lazy: bool
self, df: PandasDataFrame, keys: list[str], *, is_eager: bool, is_lazy: bool
) -> None:
self._df = df
self._keys = list(keys)
Expand All @@ -36,7 +36,7 @@ def agg(
self,
*aggs: IntoExpr | Iterable[IntoExpr],
**named_aggs: IntoExpr,
) -> PdxDataFrame:
) -> PandasDataFrame:
df = self._df._dataframe
exprs = parse_into_exprs(
self._df._implementation,
Expand Down Expand Up @@ -77,10 +77,10 @@ def agg(
self._from_dataframe,
)

def _from_dataframe(self, df: PdxDataFrame) -> PdxDataFrame:
from narwhals.pandas_like.dataframe import PdxDataFrame
def _from_dataframe(self, df: PandasDataFrame) -> PandasDataFrame:
from narwhals.pandas_like.dataframe import PandasDataFrame

return PdxDataFrame(
return PandasDataFrame(
df,
implementation=self._df._implementation,
is_eager=self._is_eager,
Expand All @@ -93,8 +93,8 @@ def agg_pandas(
exprs: list[Expr],
keys: list[str],
output_names: list[str],
from_dataframe: Callable[[Any], PdxDataFrame],
) -> PdxDataFrame:
from_dataframe: Callable[[Any], PandasDataFrame],
) -> PandasDataFrame:
"""
This should be the fastpath, but cuDF is too far behind to use it.
Expand Down Expand Up @@ -146,8 +146,8 @@ def agg_generic( # noqa: PLR0913
group_by_keys: list[str],
output_names: list[str],
implementation: str,
from_dataframe: Callable[[Any], PdxDataFrame],
) -> PdxDataFrame:
from_dataframe: Callable[[Any], PandasDataFrame],
) -> PandasDataFrame:
dfs: list[Any] = []
to_remove: list[int] = []
for i, expr in enumerate(exprs):
Expand Down
30 changes: 15 additions & 15 deletions narwhals/pandas_like/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@
from typing import Iterable

from narwhals.pandas_like import dtypes
from narwhals.pandas_like.dataframe import PdxDataFrame
from narwhals.pandas_like.dataframe import PandasDataFrame
from narwhals.pandas_like.expr import Expr
from narwhals.pandas_like.series import PandasLikeSeries
from narwhals.pandas_like.series import PandasSeries
from narwhals.pandas_like.utils import horizontal_concat
from narwhals.pandas_like.utils import parse_into_exprs
from narwhals.pandas_like.utils import series_from_iterable
Expand All @@ -31,13 +31,13 @@ class Namespace(NamespaceProtocol):
Boolean = dtypes.Boolean
String = dtypes.String

def Series(self, name: str, data: list[Any]) -> PandasLikeSeries: # noqa: N802
from narwhals.pandas_like.series import PandasLikeSeries
def Series(self, name: str, data: list[Any]) -> PandasSeries: # noqa: N802
from narwhals.pandas_like.series import PandasSeries

if self._implementation == "pandas":
import pandas as pd

return PandasLikeSeries(
return PandasSeries(
pd.Series(name=name, data=data), implementation=self._implementation
)
raise NotImplementedError
Expand All @@ -48,7 +48,7 @@ def __init__(self, implementation: str) -> None:

def _create_expr_from_callable( # noqa: PLR0913
self,
func: Callable[[PdxDataFrame], list[PandasLikeSeries]],
func: Callable[[PandasDataFrame], list[PandasSeries]],
*,
depth: int,
function_name: str,
Expand All @@ -65,9 +65,9 @@ def _create_expr_from_callable( # noqa: PLR0913
)

def _create_series_from_scalar(
self, value: Any, series: PandasLikeSeries
) -> PandasLikeSeries:
return PandasLikeSeries(
self, value: Any, series: PandasSeries
) -> PandasSeries:
return PandasSeries(
series_from_iterable(
[value],
name=series.series.name,
Expand All @@ -77,7 +77,7 @@ def _create_series_from_scalar(
implementation=self._implementation,
)

def _create_expr_from_series(self, series: PandasLikeSeries) -> Expr:
def _create_expr_from_series(self, series: PandasSeries) -> Expr:
return Expr(
lambda _df: [series],
depth=0,
Expand All @@ -96,7 +96,7 @@ def col(self, *column_names: str | Iterable[str]) -> Expr:
def all(self) -> Expr:
return Expr(
lambda df: [
PandasLikeSeries(
PandasSeries(
df._dataframe.loc[:, column_name],
implementation=self._implementation,
)
Expand Down Expand Up @@ -133,7 +133,7 @@ def min(self, *column_names: str) -> Expr:
def len(self) -> Expr:
return Expr(
lambda df: [
PandasLikeSeries(
PandasSeries(
series_from_iterable(
[len(df._dataframe)],
name="len",
Expand Down Expand Up @@ -162,10 +162,10 @@ def any_horizontal(self, *exprs: IntoExpr | Iterable[IntoExpr]) -> Expr:

def concat(
self,
items: Iterable[PdxDataFrame], # type: ignore[override]
items: Iterable[PandasDataFrame], # type: ignore[override]
*,
how: str = "vertical",
) -> PdxDataFrame:
) -> PandasDataFrame:
dfs: list[Any] = []
kind: Any = set()
for df in items:
Expand All @@ -177,7 +177,7 @@ def concat(
if how != "horizontal":
msg = "Only horizontal concatenation is supported for now"
raise TypeError(msg)
return PdxDataFrame(
return PandasDataFrame(
horizontal_concat(dfs, implementation=self._implementation),
implementation=self._implementation,
# TODO (incorrect!)
Expand Down
Loading

0 comments on commit f8cb4a5

Please sign in to comment.