refactor(expr-ir): Add _guards.py (#3092)

Author: dangotbanned

narwhals/_plan/_guards.py

@@ -0,0 +1,103 @@
+"""Common type guards, mostly with inline imports."""
+
+from __future__ import annotations
+
+import datetime as dt
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from narwhals._utils import _hasattr_static
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeIs
+
+    from narwhals._plan import expr
+    from narwhals._plan.dummy import Expr, Series
+    from narwhals._plan.protocols import CompliantSeries
+    from narwhals._plan.typing import NativeSeriesT, Seq
+    from narwhals.typing import NonNestedLiteral
+
+    T = TypeVar("T")
+
+_NON_NESTED_LITERAL_TPS = (
+    int,
+    float,
+    str,
+    dt.date,
+    dt.time,
+    dt.timedelta,
+    bytes,
+    Decimal,
+)
+
+
+def _dummy(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
+    from narwhals._plan import dummy
+
+    return dummy
+
+
+def _expr(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
+    from narwhals._plan import expr
+
+    return expr
+
+
+def is_non_nested_literal(obj: Any) -> TypeIs[NonNestedLiteral]:
+    return obj is None or isinstance(obj, _NON_NESTED_LITERAL_TPS)
+
+
+def is_expr(obj: Any) -> TypeIs[Expr]:
+    return isinstance(obj, _dummy().Expr)
+
+
+def is_column(obj: Any) -> TypeIs[Expr]:
+    """Indicate if the given object is a basic/unaliased column."""
+    return is_expr(obj) and obj.meta.is_column()
+
+
+def is_series(obj: Series[NativeSeriesT] | Any) -> TypeIs[Series[NativeSeriesT]]:
+    return isinstance(obj, _dummy().Series)
+
+
+def is_compliant_series(
+    obj: CompliantSeries[NativeSeriesT] | Any,
+) -> TypeIs[CompliantSeries[NativeSeriesT]]:
+    return _hasattr_static(obj, "__narwhals_series__")
+
+
+def is_iterable_reject(obj: Any) -> TypeIs[str | bytes | Series | CompliantSeries]:
+    return isinstance(obj, (str, bytes, _dummy().Series)) or is_compliant_series(obj)
+
+
+def is_window_expr(obj: Any) -> TypeIs[expr.WindowExpr]:
+    return isinstance(obj, _expr().WindowExpr)
+
+
+def is_function_expr(obj: Any) -> TypeIs[expr.FunctionExpr[Any]]:
+    return isinstance(obj, _expr().FunctionExpr)
+
+
+def is_binary_expr(obj: Any) -> TypeIs[expr.BinaryExpr]:
+    return isinstance(obj, _expr().BinaryExpr)
+
+
+def is_agg_expr(obj: Any) -> TypeIs[expr.AggExpr]:
+    return isinstance(obj, _expr().AggExpr)
+
+
+def is_aggregation(obj: Any) -> TypeIs[expr.AggExpr | expr.FunctionExpr[Any]]:
+    """Superset of `ExprIR.is_scalar`, excludes literals & len."""
+    return is_agg_expr(obj) or (is_function_expr(obj) and obj.is_scalar)
+
+
+def is_literal(obj: Any) -> TypeIs[expr.Literal[Any]]:
+    return isinstance(obj, _expr().Literal)
+
+
+def is_horizontal_reduction(obj: Any) -> TypeIs[expr.FunctionExpr[Any]]:
+    return is_function_expr(obj) and obj.options.is_input_wildcard_expansion()
+
+
+def is_tuple_of(obj: Any, tp: type[T]) -> TypeIs[Seq[T]]:
+    return bool(isinstance(obj, tuple) and obj and isinstance(obj[0], tp))

narwhals/_plan/arrow/namespace.py

@@ -7,9 +7,10 @@
 import pyarrow.compute as pc  # ignore-banned-import
 
 from narwhals._arrow.utils import narwhals_to_native_dtype
+from narwhals._plan._guards import is_tuple_of
 from narwhals._plan.arrow import functions as fn
 from narwhals._plan.arrow.functions import lit
-from narwhals._plan.common import collect, is_tuple_of
+from narwhals._plan.common import collect
 from narwhals._plan.literal import is_literal_scalar
 from narwhals._plan.protocols import EagerNamespace
 from narwhals._utils import Version

narwhals/_plan/common.py

@@ -8,6 +8,7 @@
 from operator import attrgetter
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast, overload
 
+from narwhals._plan._guards import is_function_expr, is_iterable_reject, is_literal
 from narwhals._plan._immutable import Immutable
 from narwhals._plan.options import ExprIROptions, FEOptions, FunctionOptions
 from narwhals._plan.typing import (
@@ -19,33 +20,22 @@
     IRNamespaceT,
     MapIR,
     NamedOrExprIRT,
-    NativeSeriesT,
     NonNestedDTypeT,
     Seq,
 )
-from narwhals._utils import _hasattr_static
 from narwhals.dtypes import DType
 from narwhals.utils import Version
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
     from typing import Any, Callable
 
-    from typing_extensions import Self, TypeAlias, TypeIs
-
-    from narwhals._plan import expr
-    from narwhals._plan.dummy import Expr, Selector, Series
-    from narwhals._plan.expr import (
-        AggExpr,
-        Alias,
-        BinaryExpr,
-        Cast,
-        Column,
-        FunctionExpr,
-        WindowExpr,
-    )
+    from typing_extensions import Self, TypeAlias
+
+    from narwhals._plan.dummy import Expr, Selector
+    from narwhals._plan.expr import Alias, Cast, Column, FunctionExpr
     from narwhals._plan.meta import IRMetaNamespace
-    from narwhals._plan.protocols import CompliantSeries, Ctx, FrameT_contra, R_co
+    from narwhals._plan.protocols import Ctx, FrameT_contra, R_co
     from narwhals.typing import NonNestedDType, NonNestedLiteral
 
 
@@ -453,97 +443,6 @@ class HorizontalFunction(
 ): ...
 
 
-_NON_NESTED_LITERAL_TPS = (
-    int,
-    float,
-    str,
-    dt.date,
-    dt.time,
-    dt.timedelta,
-    bytes,
-    Decimal,
-)
-
-
-def is_non_nested_literal(obj: Any) -> TypeIs[NonNestedLiteral]:
-    return obj is None or isinstance(obj, _NON_NESTED_LITERAL_TPS)
-
-
-def is_expr(obj: Any) -> TypeIs[Expr]:
-    from narwhals._plan.dummy import Expr
-
-    return isinstance(obj, Expr)
-
-
-def is_column(obj: Any) -> TypeIs[Expr]:
-    """Indicate if the given object is a basic/unaliased column.
-
-    https://github.com/pola-rs/polars/blob/a3d6a3a7863b4d42e720a05df69ff6b6f5fc551f/py-polars/polars/_utils/various.py#L164-L168.
-    """
-    return is_expr(obj) and obj.meta.is_column()
-
-
-def is_series(obj: Series[NativeSeriesT] | Any) -> TypeIs[Series[NativeSeriesT]]:
-    from narwhals._plan.dummy import Series
-
-    return isinstance(obj, Series)
-
-
-def is_compliant_series(
-    obj: CompliantSeries[NativeSeriesT] | Any,
-) -> TypeIs[CompliantSeries[NativeSeriesT]]:
-    return _hasattr_static(obj, "__narwhals_series__")
-
-
-def is_iterable_reject(obj: Any) -> TypeIs[str | bytes | Series | CompliantSeries]:
-    from narwhals._plan.dummy import Series
-
-    return isinstance(obj, (str, bytes, Series)) or is_compliant_series(obj)
-
-
-def is_window_expr(obj: Any) -> TypeIs[WindowExpr]:
-    from narwhals._plan.expr import WindowExpr
-
-    return isinstance(obj, WindowExpr)
-
-
-def is_function_expr(obj: Any) -> TypeIs[FunctionExpr[Any]]:
-    from narwhals._plan.expr import FunctionExpr
-
-    return isinstance(obj, FunctionExpr)
-
-
-def is_binary_expr(obj: Any) -> TypeIs[BinaryExpr]:
-    from narwhals._plan.expr import BinaryExpr
-
-    return isinstance(obj, BinaryExpr)
-
-
-def is_agg_expr(obj: Any) -> TypeIs[AggExpr]:
-    from narwhals._plan.expr import AggExpr
-
-    return isinstance(obj, AggExpr)
-
-
-def is_aggregation(obj: Any) -> TypeIs[AggExpr | FunctionExpr[Any]]:
-    """Superset of `ExprIR.is_scalar`, excludes literals & len."""
-    return is_agg_expr(obj) or (is_function_expr(obj) and obj.is_scalar)
-
-
-def is_literal(obj: Any) -> TypeIs[expr.Literal[Any]]:
-    from narwhals._plan import expr
-
-    return isinstance(obj, expr.Literal)
-
-
-def is_horizontal_reduction(obj: FunctionExpr[Any] | Any) -> TypeIs[FunctionExpr[Any]]:
-    return is_function_expr(obj) and obj.options.is_input_wildcard_expansion()
-
-
-def is_tuple_of(obj: Any, tp: type[T]) -> TypeIs[Seq[T]]:
-    return bool(isinstance(obj, tuple) and obj and isinstance(obj[0], tp))
-
-
 def py_to_narwhals_dtype(obj: NonNestedLiteral, version: Version = Version.MAIN) -> DType:
     dtypes = version.dtypes
     mapping: dict[type[NonNestedLiteral], type[NonNestedDType]] = {

narwhals/_plan/demo.py

@@ -3,13 +3,8 @@
 import builtins
 import typing as t
 
-from narwhals._plan import boolean, expr, expr_parsing as parse, functions as F
-from narwhals._plan.common import (
-    into_dtype,
-    is_non_nested_literal,
-    is_series,
-    py_to_narwhals_dtype,
-)
+from narwhals._plan import _guards, boolean, expr, expr_parsing as parse, functions as F
+from narwhals._plan.common import into_dtype, py_to_narwhals_dtype
 from narwhals._plan.expr import All, Len
 from narwhals._plan.literal import ScalarLiteral, SeriesLiteral
 from narwhals._plan.ranges import IntRange
@@ -39,9 +34,9 @@ def nth(*indices: int | t.Sequence[int]) -> Expr:
 def lit(
     value: NonNestedLiteral | Series[NativeSeriesT], dtype: IntoDType | None = None
 ) -> Expr:
-    if is_series(value):
+    if _guards.is_series(value):
         return SeriesLiteral(value=value).to_literal().to_narwhals()
-    if not is_non_nested_literal(value):
+    if not _guards.is_non_nested_literal(value):
         msg = f"{type(value).__name__!r} is not supported in `nw.lit`, got: {value!r}."
         raise TypeError(msg)
     if dtype is None:

narwhals/_plan/dummy.py

@@ -15,7 +15,8 @@
     functions as F,
     operators as ops,
 )
-from narwhals._plan.common import NamedIR, into_dtype, is_column, is_expr, is_series
+from narwhals._plan._guards import is_column, is_expr, is_series
+from narwhals._plan.common import NamedIR, into_dtype
 from narwhals._plan.contexts import ExprContext
 from narwhals._plan.options import (
     EWMOptions,

narwhals/_plan/expr_expansion.py

@@ -43,8 +43,9 @@
 from typing import TYPE_CHECKING
 
 from narwhals._plan import common, meta
+from narwhals._plan._guards import is_horizontal_reduction
 from narwhals._plan._immutable import Immutable
-from narwhals._plan.common import ExprIR, NamedIR, SelectorIR, is_horizontal_reduction
+from narwhals._plan.common import ExprIR, NamedIR, SelectorIR
 from narwhals._plan.exceptions import (
     column_index_error,
     column_not_found_error,

narwhals/_plan/expr_parsing.py

@@ -6,7 +6,7 @@
 from itertools import chain
 from typing import TYPE_CHECKING, TypeVar
 
-from narwhals._plan.common import is_expr, is_iterable_reject
+from narwhals._plan._guards import is_expr, is_iterable_reject
 from narwhals._plan.exceptions import (
     invalid_into_expr_error,
     is_iterable_pandas_error,

narwhals/_plan/expr_rewrites.py

@@ -5,15 +5,13 @@
 from typing import TYPE_CHECKING
 
 from narwhals._plan import expr_parsing as parse
-from narwhals._plan.common import (
-    NamedIR,
+from narwhals._plan._guards import (
     is_aggregation,
     is_binary_expr,
     is_function_expr,
     is_window_expr,
-    map_ir,
-    replace,
 )
+from narwhals._plan.common import NamedIR, map_ir, replace
 from narwhals._plan.expr_expansion import into_named_irs, prepare_projection
 
 if TYPE_CHECKING:

narwhals/_plan/literal.py

@@ -2,8 +2,8 @@
 
 from typing import TYPE_CHECKING, Any, Generic
 
+from narwhals._plan._guards import is_literal
 from narwhals._plan._immutable import Immutable
-from narwhals._plan.common import is_literal
 from narwhals._plan.typing import LiteralT, NativeSeriesT, NonNestedLiteralT
 
 if TYPE_CHECKING:

narwhals/_plan/operators.py

@@ -3,8 +3,8 @@
 import operator as op
 from typing import TYPE_CHECKING
 
+from narwhals._plan._guards import is_function_expr
 from narwhals._plan._immutable import Immutable
-from narwhals._plan.common import is_function_expr
 from narwhals._plan.exceptions import (
     binary_expr_length_changing_error,
     binary_expr_multi_output_error,