|
4 | 4 | from importlib import import_module |
5 | 5 | from typing import TYPE_CHECKING |
6 | 6 | from typing import Any |
| 7 | +from typing import Generic |
7 | 8 | from typing import Iterator |
8 | 9 | from typing import Literal |
9 | 10 | from typing import Sequence |
| 11 | +from typing import TypeVar |
10 | 12 | from typing import cast |
11 | 13 |
|
12 | 14 | from narwhals._spark_like.utils import evaluate_exprs |
|
41 | 43 | from narwhals.utils import Version |
42 | 44 |
|
43 | 45 | SQLFrameDataFrame: TypeAlias = _SQLFrameDataFrame[Any, Any, Any, Any, Any] |
44 | | - _NativeDataFrame: TypeAlias = "DataFrame | SQLFrameDataFrame" |
| 46 | + _NativeDataFrame: TypeAlias = "DataFrame | SQLFrameDataFrame" # noqa: PYI047 |
45 | 47 |
|
| 48 | +FrameT = TypeVar("FrameT", "DataFrame", "SQLFrameDataFrame") |
46 | 49 | Incomplete: TypeAlias = Any # pragma: no cover |
47 | 50 | """Marker for working code that fails type checking.""" |
48 | 51 |
|
49 | 52 |
|
50 | | -class SparkLikeLazyFrame(CompliantLazyFrame): |
| 53 | +class SparkLikeLazyFrame(CompliantLazyFrame, Generic[FrameT]): |
51 | 54 | def __init__( |
52 | 55 | self: Self, |
53 | | - native_dataframe: _NativeDataFrame, |
| 56 | + native_dataframe: FrameT, |
54 | 57 | *, |
55 | 58 | backend_version: tuple[int, ...], |
56 | 59 | version: Version, |
@@ -144,7 +147,7 @@ def _change_version(self: Self, version: Version) -> Self: |
144 | 147 | implementation=self._implementation, |
145 | 148 | ) |
146 | 149 |
|
147 | | - def _from_native_frame(self: Self, df: DataFrame) -> Self: |
| 150 | + def _from_native_frame(self: Self, df: FrameT) -> Self: |
148 | 151 | return self.__class__( |
149 | 152 | df, |
150 | 153 | backend_version=self._backend_version, |
|
0 commit comments