1313
1414if TYPE_CHECKING :
1515 from narwhals ._typing import Arrow , IntoBackend , Pandas , Polars
16+ from narwhals .typing import DataFrameT , LazyFrameT
1617
1718GUARANTEES_ROW_ORDER = {
1819 Implementation .PANDAS ,
2526
2627
2728def assert_frame_equal (
28- left : DataFrame [ Any ] | LazyFrame [ Any ] ,
29- right : DataFrame [ Any ] | LazyFrame [ Any ] ,
29+ left : DataFrameT | LazyFrameT ,
30+ right : DataFrameT | LazyFrameT ,
3031 * ,
3132 check_row_order : bool = True ,
3233 check_column_order : bool = True ,
@@ -144,8 +145,8 @@ def assert_frame_equal(
144145
145146
146147def _check_correct_input_type ( # noqa: RET503
147- left : DataFrame [ Any ] | LazyFrame [ Any ] ,
148- right : DataFrame [ Any ] | LazyFrame [ Any ] ,
148+ left : DataFrameT | LazyFrameT ,
149+ right : DataFrameT | LazyFrameT ,
149150 backend : IntoBackend [Polars | Pandas | Arrow ] | None ,
150151) -> tuple [DataFrame [Any ], DataFrame [Any ]]:
151152 # Adapted from https://github.com/pola-rs/polars/blob/afdbf3056d1228cf493901e45f536b0905cec8ea/py-polars/src/polars/testing/asserts/frame.py#L15-L17
@@ -164,8 +165,8 @@ def _check_correct_input_type( # noqa: RET503
164165
165166
166167def _assert_dataframe_equal (
167- left : DataFrame [ Any ] ,
168- right : DataFrame [ Any ] ,
168+ left : DataFrameT ,
169+ right : DataFrameT ,
169170 impl : Implementation ,
170171 * ,
171172 check_row_order : bool ,
@@ -231,11 +232,7 @@ def _assert_dataframe_equal(
231232
232233
233234def _check_schema_equal (
234- left : DataFrame [Any ],
235- right : DataFrame [Any ],
236- * ,
237- check_dtypes : bool ,
238- check_column_order : bool ,
235+ left : DataFrameT , right : DataFrameT , * , check_dtypes : bool , check_column_order : bool
239236) -> None :
240237 """Compares DataFrame schema based on specified criteria.
241238
0 commit comments