Skip to content

Commit

Permalink
oooh!
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed Mar 11, 2024
1 parent 5df15be commit e6d5c67
Show file tree
Hide file tree
Showing 5 changed files with 100 additions and 3 deletions.
4 changes: 3 additions & 1 deletion narwhals/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@
from narwhals.containers import is_pandas
from narwhals.containers import is_polars
from narwhals.containers import is_series
from narwhals.expressions import col
from narwhals.dataframe import NarwhalsFrame
from narwhals.expression import col
from narwhals.translate import get_namespace
from narwhals.translate import to_native
from narwhals.translate import translate_any
Expand All @@ -24,4 +25,5 @@
"get_namespace",
"to_native",
"col",
"NarwhalsFrame",
]
92 changes: 92 additions & 0 deletions narwhals/dataframe.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
from __future__ import annotations

from narwhals.pandas_like.utils import evaluate_into_exprs
from narwhals.translate import get_pandas
from narwhals.translate import get_polars


def extract_native(obj: Any, implementation) -> Any:
from narwhals.expression import NarwhalsExpr

# if isinstance(obj, NarwhalsExpr):
# return obj._call(pl.col)
if isinstance(obj, NarwhalsExpr):
if implementation == "polars":
import polars as pl

return obj._call(pl.col)
# if isinstance(obj, DType):
# return obj._dtype
if isinstance(obj, NarwhalsFrame):
return obj._dataframe
# if isinstance(obj, PolarsSeries):
# return obj._series
return obj


class NarwhalsFrame:
def __init__(
self, df, *, is_eager=False, is_lazy=False, implementation: str | None = None
):
self._is_eager = is_eager
self._is_lazy = is_lazy
if implementation is not None:
self._dataframe = df
self._implementation = implementation
return
if (pl := get_polars()) is not None:
if isinstance(df, pl.DataFrame):
if is_lazy:
raise ValueError(
"can't instantiate with `is_lazy` if you pass a polars DataFrame"
)
self._dataframe = df
self._implementation = "polars"
return
elif isinstance(df, pl.LazyFrame):
if is_eager:
raise ValueError(
"can't instantiate with `is_eager` if you pass a polars LazyFrame"
)
self._dataframe = df
self._implementation = "polars"
return
if (pd := get_pandas()) is not None and isinstance(df, pd.DataFrame):
self._dataframe = df
self._implementation = "pandas"
return
raise TypeError(
f"Expected pandas or Polars dataframe or lazyframe, got: {type(df)}"
)

def _from_dataframe(self, df: Any) -> Self:
# construct, preserving properties
return self.__class__(
df,
is_eager=self._is_eager,
is_lazy=self._is_lazy,
implementation=self._implementation,
)

def _extract_native(self, obj):
return extract_native(obj, implementation=self._implementation)

def with_columns(
self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr
) -> Self:
if self._implementation == "polars":
return self._from_dataframe(
self._dataframe.with_columns(
*[self._extract_native(v) for v in exprs],
**{
key: self._extract_native(value)
for key, value in named_exprs.items()
},
)
)
elif self._implementation == "pandas":
new_series = evaluate_into_exprs(self, *exprs, **named_exprs)
df = self._dataframe.assign(
**{series.name: series._series for series in new_series}
)
return self._from_dataframe(df)
2 changes: 1 addition & 1 deletion narwhals/pandas_like/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def parse_into_exprs(


def parse_into_expr(implementation: str, into_expr: IntoExpr) -> Expr:
from narwhals.expressions import NarwhalsExpr
from narwhals.expression import NarwhalsExpr
from narwhals.pandas_like.expr import Expr
from narwhals.pandas_like.namespace import Namespace
from narwhals.pandas_like.series import PandasSeries
Expand Down
2 changes: 1 addition & 1 deletion narwhals/polars.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@


def extract_native(obj: Any) -> Any:
from narwhals.expressions import NarwhalsExpr
from narwhals.expression import NarwhalsExpr

if isinstance(obj, NarwhalsExpr):
return obj._call(pl.col)
Expand Down
3 changes: 3 additions & 0 deletions narwhals/translate.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,11 +175,14 @@ def get_namespace(obj: Any) -> Namespace:


def to_native(obj: Any) -> Any:
from narwhals.dataframe import NarwhalsFrame
from narwhals.pandas_like.dataframe import PandasDataFrame
from narwhals.pandas_like.series import PandasSeries
from narwhals.polars import PolarsDataFrame
from narwhals.polars import PolarsSeries

if isinstance(obj, NarwhalsFrame):
return obj._dataframe
if isinstance(obj, PandasDataFrame):
return obj._dataframe
if isinstance(obj, PandasSeries):
Expand Down

0 comments on commit e6d5c67

Please sign in to comment.