simplify dask imports

Author: MarcoGorelli

narwhals/_dask/dataframe.py

@@ -7,6 +7,8 @@
 from typing import Literal
 from typing import Sequence
 
+import dask.dataframe as dd
+
 from narwhals._dask.utils import add_row_index
 from narwhals._dask.utils import parse_exprs_and_named_exprs
 from narwhals._pandas_like.utils import native_to_narwhals_dtype
@@ -23,7 +25,6 @@
 if TYPE_CHECKING:
     from types import ModuleType
 
-    import dask.dataframe as dd
     from typing_extensions import Self
 
     from narwhals._dask.expr import DaskExpr
@@ -111,8 +112,6 @@ def select(
         *exprs: IntoDaskExpr,
         **named_exprs: IntoDaskExpr,
     ) -> Self:
-        import dask.dataframe as dd
-
         if exprs and all(isinstance(x, str) for x in exprs) and not named_exprs:
             # This is a simple slice => fastpath!
             return self._from_native_frame(

narwhals/_dask/expr_str.py

@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING
 
+import dask.dataframe as dd
+
 if TYPE_CHECKING:
     from typing_extensions import Self
 
@@ -93,8 +95,6 @@ def slice(self, offset: int, length: int | None = None) -> DaskExpr:
         )
 
     def to_datetime(self: Self, format: str | None) -> DaskExpr:  # noqa: A002
-        import dask.dataframe as dd
-
         return self._compliant_expr._from_call(
             lambda _input, format: dd.to_datetime(_input, format=format),
             "to_datetime",

narwhals/_dask/group_by.py

@@ -6,19 +6,20 @@
 from typing import Callable
 from typing import Sequence
 
+import dask.dataframe as dd
+
+try:
+    import dask.dataframe.dask_expr as dx
+except ModuleNotFoundError:  # pragma: no cover
+    import dask_expr as dx
+
+
 from narwhals._expression_parsing import is_simple_aggregation
 from narwhals._expression_parsing import parse_into_exprs
 from narwhals.exceptions import AnonymousExprError
 from narwhals.utils import remove_prefix
 
 if TYPE_CHECKING:
-    import dask.dataframe as dd
-
-    try:
-        import dask.dataframe.dask_expr as dx
-    except ModuleNotFoundError:
-        import dask_expr as dx
-
     import pandas as pd
 
     from narwhals._dask.dataframe import DaskLazyFrame
@@ -27,8 +28,6 @@
 
 
 def n_unique() -> dd.Aggregation:
-    import dask.dataframe as dd
-
     def chunk(s: pd.core.groupby.generic.SeriesGroupBy) -> int:
         return s.nunique(dropna=False)  # type: ignore[no-any-return]
 
@@ -49,11 +48,6 @@ def var(
 ]:
     from functools import partial
 
-    try:
-        import dask.dataframe.dask_expr as dx
-    except ModuleNotFoundError:  # pragma: no cover
-        import dask_expr as dx
-
     return partial(dx._groupby.GroupBy.var, ddof=ddof)
 
 
@@ -64,11 +58,6 @@ def std(
 ]:
     from functools import partial
 
-    try:
-        import dask.dataframe.dask_expr as dx
-    except ModuleNotFoundError:  # pragma: no cover
-        import dask_expr as dx
-
     return partial(dx._groupby.GroupBy.std, ddof=ddof)
 
 

narwhals/_dask/namespace.py

@@ -8,6 +8,9 @@
 from typing import Sequence
 from typing import cast
 
+import dask.dataframe as dd
+import pandas as pd
+
 from narwhals._dask.dataframe import DaskLazyFrame
 from narwhals._dask.expr import DaskExpr
 from narwhals._dask.selectors import DaskSelectorNamespace
@@ -69,9 +72,6 @@ def nth(self, *column_indices: int) -> DaskExpr:
         )
 
     def lit(self, value: Any, dtype: DType | None) -> DaskExpr:
-        import dask.dataframe as dd
-        import pandas as pd
-
         def func(df: DaskLazyFrame) -> list[dx.Series]:
             return [
                 dd.from_pandas(
@@ -99,7 +99,6 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
         )
 
     def len(self) -> DaskExpr:
-        import dask.dataframe as dd
         import pandas as pd
 
         def func(df: DaskLazyFrame) -> list[dx.Series]:
@@ -188,8 +187,6 @@ def concat(
         *,
         how: Literal["horizontal", "vertical", "diagonal"],
     ) -> DaskLazyFrame:
-        import dask.dataframe as dd
-
         if len(list(items)) == 0:
             msg = "No items to concatenate"  # pragma: no cover
             raise AssertionError(msg)
@@ -265,8 +262,6 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
         )
 
     def min_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr:
-        import dask.dataframe as dd
-
         parsed_exprs = parse_into_exprs(*exprs, namespace=self)
 
         def func(df: DaskLazyFrame) -> list[dx.Series]:
@@ -287,8 +282,6 @@ def func(df: DaskLazyFrame) -> list[dx.Series]:
         )
 
     def max_horizontal(self, *exprs: IntoDaskExpr) -> DaskExpr:
-        import dask.dataframe as dd
-
         parsed_exprs = parse_into_exprs(*exprs, namespace=self)
 
         def func(df: DaskLazyFrame) -> list[dx.Series]: