wip operator

Author: ym-pett

narwhals/_sql/expr.py

@@ -15,7 +15,7 @@
     combine_alias_output_names,
     combine_evaluate_output_names,
 )
-from narwhals._sql.typing import SQLLazyFrameT
+from narwhals._sql.typing import SQLLazyFrameT, NativeSQLExprT
 from narwhals._utils import Implementation, Version, not_implemented
 
 if TYPE_CHECKING:
@@ -249,7 +249,7 @@ def _rolling_window_func(
             end = 0
 
         def func(
-            df: SQLLazyFrameT, inputs: WindowInputs[NativeExprT]
+            df: SQLLazyFrameT, inputs: WindowInputs[NativeSQLExprT]
         ) -> Sequence[NativeExprT]:
             if func_name in {"sum", "mean"}:
                 func_: str = func_name
@@ -496,11 +496,11 @@ def round(self, decimals: int) -> Self:
         return self._with_elementwise(
             lambda expr: self._function("round", expr, self._lit(decimals))
         )
-
+    # WIP: trying new NativeSQLExprT 
     def sqrt(self) -> Self:
-        def _sqrt(expr: NativeExprT) -> NativeExprT:
+        def _sqrt(expr: NativeSQLExprT) -> NativeSQLExprT:
             return self._when(
-                expr < self._lit(0),  # type: ignore[operator]
+                expr < self._lit(0),  
                 self._lit(float("nan")),
                 self._function("sqrt", expr),
             )

narwhals/_sql/typing.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar, Self
 
 from narwhals._compliant.expr import NativeExpr
 
@@ -9,32 +9,33 @@
     from narwhals._sql.expr import SQLExpr
     from narwhals.dtypes import Boolean
 
+    # TODO: @mp, understand why these are here & if we need one for NativeSQLExprT
     SQLExprAny = SQLExpr[Any, Any]
     SQLLazyFrameAny = SQLLazyFrame[Any, Any, Any]
 
 SQLExprT = TypeVar("SQLExprT", bound="SQLExprAny")
 SQLExprT_contra = TypeVar("SQLExprT_contra", bound="SQLExprAny", contravariant=True)
 SQLLazyFrameT = TypeVar("SQLLazyFrameT", bound="SQLLazyFrameAny")
-
+NativeSQLExprT = TypeVar("NativeSQLExprT", bound="NativeSQLExpr")
 
 class NativeSQLExpr(NativeExpr):
-    # TODO @mp: fix input type for all these!
-    def __gt__(self, value: float) -> Boolean: ...
+    # both Self because we're comparing an expression with an expression? 
+    def __gt__(self, value: Self) -> Self: ...
 
-    def __lt__(self, value: float) -> Boolean: ...
+    def __lt__(self, value: Self) -> Self: ...
 
-    def __ge__(self, value: float) -> Boolean: ...
+    def __ge__(self, value: Self) -> Self: ...
 
-    def __le__(self, value: float) -> Boolean: ...
+    def __le__(self, value: Self) -> Self: ...
 
-    def __eq__(self, value: float) -> Boolean: ...
+    def __eq__(self, value: Self) -> Self: ...
 
-    def __ne__(self, value: float) -> Boolean: ...
+    def __ne__(self, value: Self) -> Self: ...
     # do we want any more of the arithmetic methods? I wasn't sure between lefthan & righthand methods..
-    def __sub__(self, value: float) -> Any: ...
+    def __sub__(self, value: Self) -> Self: ...
 
-    def __add__(self, value: float) -> Any: ...
+    def __add__(self, value: Self) -> Self: ...
 
-    def __truediv__(self, value: float) -> Any: ...
+    def __truediv__(self, value: Self) -> Self: ...
 
-    def __mul__(self, value: float) -> Any: ...
+    def __mul__(self, value: Self) -> Self: ...