fix(typing): factor out imports, avoid redef

- `mypy` just needed the `else` block
- but this is much more maintainable

Author: dangotbanned

narwhals/_spark_like/dataframe.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
 import warnings
-from importlib import import_module
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterator
 from typing import Literal
 from typing import Sequence
 
 from narwhals._spark_like.utils import evaluate_exprs
+from narwhals._spark_like.utils import import_functions
+from narwhals._spark_like.utils import import_native_dtypes
+from narwhals._spark_like.utils import import_window
 from narwhals._spark_like.utils import native_to_narwhals_dtype
 from narwhals.exceptions import InvalidOperationError
 from narwhals.typing import CompliantDataFrame
@@ -70,51 +72,30 @@ def _F(self: Self):  # type: ignore[no-untyped-def] # noqa: ANN202, N802
 
             return functions
         else:
-            if self._implementation is Implementation.SQLFRAME:
-                from sqlframe.base.session import _BaseSession
-
-                return import_module(
-                    f"sqlframe.{_BaseSession().execution_dialect_name}.functions"
-                )
-
-            from pyspark.sql import functions
-
-            return functions
+            return import_functions(self._implementation)
 
     @property
     def _native_dtypes(self: Self):  # type: ignore[no-untyped-def] # noqa: ANN202
         if TYPE_CHECKING:
             from sqlframe.base import types
 
             return types
-
-        if self._implementation is Implementation.SQLFRAME:
-            from sqlframe.base.session import _BaseSession
-
-            return import_module(
-                f"sqlframe.{_BaseSession().execution_dialect_name}.types"
-            )
-
-        from pyspark.sql import types  # type: ignore[no-redef]
-
-        return types
+        else:
+            return import_native_dtypes(self._implementation)
 
     @property
     def _Window(self: Self) -> type[Window]:  # noqa: N802
-        if self._implementation is Implementation.SQLFRAME:
-            from sqlframe.base.session import _BaseSession
-
-            _window = import_module(
-                f"sqlframe.{_BaseSession().execution_dialect_name}.window"
-            )
-            return _window.Window
-
-        from pyspark.sql import Window
+        if TYPE_CHECKING:
+            from sqlframe.base.window import Window
 
-        return Window
+            return Window
+        else:
+            return import_window(self._implementation)
 
     @property
     def _session(self: Self) -> SQLFrameSession:
+        if TYPE_CHECKING:
+            return self._native_frame.session
         if self._implementation is Implementation.SQLFRAME:
             return self._native_frame.session
 

narwhals/_spark_like/expr.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import operator
-from importlib import import_module
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
@@ -16,6 +15,9 @@
 from narwhals._spark_like.expr_name import SparkLikeExprNameNamespace
 from narwhals._spark_like.expr_str import SparkLikeExprStringNamespace
 from narwhals._spark_like.expr_struct import SparkLikeExprStructNamespace
+from narwhals._spark_like.utils import import_functions
+from narwhals._spark_like.utils import import_native_dtypes
+from narwhals._spark_like.utils import import_window
 from narwhals._spark_like.utils import maybe_evaluate_expr
 from narwhals._spark_like.utils import narwhals_to_native_dtype
 from narwhals.dependencies import get_pyspark
@@ -87,47 +89,26 @@ def _F(self: Self):  # type: ignore[no-untyped-def] # noqa: ANN202, N802
             from sqlframe.base import functions
 
             return functions
-        if self._implementation is Implementation.SQLFRAME:
-            from sqlframe.base.session import _BaseSession
-
-            return import_module(
-                f"sqlframe.{_BaseSession().execution_dialect_name}.functions"
-            )
-
-        from pyspark.sql import functions  # type: ignore[no-redef]
-
-        return functions
+        else:
+            return import_functions(self._implementation)
 
     @property
     def _native_dtypes(self: Self):  # type: ignore[no-untyped-def] # noqa: ANN202
         if TYPE_CHECKING:
             from sqlframe.base import types
 
             return types
-        if self._implementation is Implementation.SQLFRAME:
-            from sqlframe.base.session import _BaseSession
-
-            return import_module(
-                f"sqlframe.{_BaseSession().execution_dialect_name}.types"
-            )
-
-        from pyspark.sql import types  # type: ignore[no-redef]
-
-        return types
+        else:
+            return import_native_dtypes(self._implementation)
 
     @property
     def _Window(self: Self) -> type[Window]:  # noqa: N802
-        if self._implementation is Implementation.SQLFRAME:
-            from sqlframe.base.session import _BaseSession
-
-            _window = import_module(
-                f"sqlframe.{_BaseSession().execution_dialect_name}.window"
-            )
-            return _window.Window
-
-        from pyspark.sql import Window
+        if TYPE_CHECKING:
+            from sqlframe.base.window import Window
 
-        return Window
+            return Window
+        else:
+            return import_window(self._implementation)
 
     def __narwhals_expr__(self: Self) -> None: ...
 

narwhals/_spark_like/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from importlib import import_module
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import cast
@@ -246,3 +247,40 @@ def _var_pyspark(col: Any, ddof: int, /) -> Column:
     from pyspark.pandas.spark.functions import var
 
     return cast("Column", var(col, ddof=ddof))
+
+
+def import_functions(implementation: Implementation, /) -> ModuleType:
+    if implementation is Implementation.SQLFRAME:
+        from sqlframe.base.session import _BaseSession
+
+        return import_module(
+            f"sqlframe.{_BaseSession().execution_dialect_name}.functions"
+        )
+
+    from pyspark.sql import functions
+
+    return functions
+
+
+def import_native_dtypes(implementation: Implementation, /) -> ModuleType:
+    if implementation is Implementation.SQLFRAME:
+        from sqlframe.base.session import _BaseSession
+
+        return import_module(f"sqlframe.{_BaseSession().execution_dialect_name}.types")
+
+    from pyspark.sql import types
+
+    return types
+
+
+def import_window(implementation: Implementation, /) -> type[Any]:
+    if implementation is Implementation.SQLFRAME:
+        from sqlframe.base.session import _BaseSession
+
+        return import_module(
+            f"sqlframe.{_BaseSession().execution_dialect_name}.window"
+        ).Window
+
+    from pyspark.sql import Window
+
+    return Window