Skip to content

Commit c4604bc

Browse files
committed
fix(typing): Generic SparkLikeLazyFrame
Drops the `pyright` errors **24** -> **9** Will close #2185
1 parent 01830c1 commit c4604bc

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

narwhals/_spark_like/dataframe.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
from importlib import import_module
55
from typing import TYPE_CHECKING
66
from typing import Any
7+
from typing import Generic
78
from typing import Iterator
89
from typing import Literal
910
from typing import Sequence
11+
from typing import TypeVar
1012
from typing import cast
1113

1214
from narwhals._spark_like.utils import evaluate_exprs
@@ -41,16 +43,17 @@
4143
from narwhals.utils import Version
4244

4345
SQLFrameDataFrame: TypeAlias = _SQLFrameDataFrame[Any, Any, Any, Any, Any]
44-
_NativeDataFrame: TypeAlias = "DataFrame | SQLFrameDataFrame"
46+
_NativeDataFrame: TypeAlias = "DataFrame | SQLFrameDataFrame" # noqa: PYI047
4547

48+
FrameT = TypeVar("FrameT", "DataFrame", "SQLFrameDataFrame")
4649
Incomplete: TypeAlias = Any # pragma: no cover
4750
"""Marker for working code that fails type checking."""
4851

4952

50-
class SparkLikeLazyFrame(CompliantLazyFrame):
53+
class SparkLikeLazyFrame(CompliantLazyFrame, Generic[FrameT]):
5154
def __init__(
5255
self: Self,
53-
native_dataframe: _NativeDataFrame,
56+
native_dataframe: FrameT,
5457
*,
5558
backend_version: tuple[int, ...],
5659
version: Version,
@@ -144,7 +147,7 @@ def _change_version(self: Self, version: Version) -> Self:
144147
implementation=self._implementation,
145148
)
146149

147-
def _from_native_frame(self: Self, df: DataFrame) -> Self:
150+
def _from_native_frame(self: Self, df: FrameT) -> Self:
148151
return self.__class__(
149152
df,
150153
backend_version=self._backend_version,

0 commit comments

Comments
 (0)