From e6d5c67ae8918c375ba44905ff10a05a95300952 Mon Sep 17 00:00:00 2001
From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com>
Date: Mon, 11 Mar 2024 19:13:40 +0000
Subject: [PATCH] oooh!

---
 narwhals/__init__.py          |  4 +-
 narwhals/dataframe.py         | 92 +++++++++++++++++++++++++++++++++++
 narwhals/pandas_like/utils.py |  2 +-
 narwhals/polars.py            |  2 +-
 narwhals/translate.py         |  3 ++
 5 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 7f57751b37..752d14a833 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -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
@@ -24,4 +25,5 @@
     "get_namespace",
     "to_native",
     "col",
+    "NarwhalsFrame",
 ]
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index e69de29bb2..ef006e8b0d 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -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)
diff --git a/narwhals/pandas_like/utils.py b/narwhals/pandas_like/utils.py
index 479fe5f738..8378b5e845 100644
--- a/narwhals/pandas_like/utils.py
+++ b/narwhals/pandas_like/utils.py
@@ -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
diff --git a/narwhals/polars.py b/narwhals/polars.py
index f11eb190e0..23d3f884b2 100644
--- a/narwhals/polars.py
+++ b/narwhals/polars.py
@@ -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)
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 0f4c7b7161..5f6d1a4124 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -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):