feat(typing): make common Expr.__dunder__ methods more type specific

ibis/expr/datatypes/value.py

@@ -10,7 +10,7 @@
 from collections.abc import Mapping, Sequence
 from functools import partial
 from operator import attrgetter
-from typing import Any
+from typing import TYPE_CHECKING, Any, Union
 
 import toolz
 from public import public
@@ -28,9 +28,56 @@
 )
 from ibis.expr.datatypes.cast import highest_precedence
 
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
+InferrableToStruct = Mapping[str, Any]
+InferrableToMap = Mapping[Any, Any]
+InferrableToArray = Union[list, tuple, set, frozenset]
+InferrableToTime = datetime.time
+InferrableToDate = datetime.date
+InferrableToTimestamp = Union[datetime.datetime, "pd.Timestamp"]
+InferrableToInterval = Union[datetime.timedelta, "pd.Timedelta"]
+InferrableToString = str | enum.Enum
+InferrableToBytes = bytes
+InferrableToFloating = float
+InferrableToInteger = int
+InferrableToDecimal = decimal.Decimal
+InferrableToBoolean = bool
+InferrableToNull = None
+InferrableToUUID = uuid.UUID
+InferrableToINET = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
+InferrableToSomething = Union["np.ndarray", "pd.Series"]
+
+InferrableToNumeric = Union[
+    InferrableToFloating, InferrableToInteger, InferrableToDecimal, InferrableToBoolean
+]
+
+Inferrable = Union[
+    InferrableToStruct,
+    InferrableToMap,
+    InferrableToArray,
+    InferrableToTime,
+    InferrableToDate,
+    InferrableToTimestamp,
+    InferrableToInterval,
+    InferrableToString,
+    InferrableToBytes,
+    InferrableToFloating,
+    InferrableToInteger,
+    InferrableToDecimal,
+    InferrableToBoolean,
+    InferrableToNull,
+    InferrableToUUID,
+    InferrableToINET,
+    InferrableToSomething,
+]
+"""All the types that can be inferred to an ibis dtype."""
+
 
 @lazy_singledispatch
-def infer(value: Any) -> dt.DataType:
+def infer(value: Inferrable) -> dt.DataType:
     """Infer the corresponding ibis dtype for a python object."""
     raise InputTypeError(
         f"Unable to infer datatype of value {value!r} with type {type(value)}"
@@ -169,22 +216,22 @@ def infer_ipaddr(
 
 
 @infer.register("numpy.generic")
-def infer_numpy_scalar(value):
+def infer_numpy_scalar(value: np.generic):
     from ibis.formats.numpy import NumpyType
 
     return NumpyType.to_ibis(value.dtype)
 
 
 @infer.register("pandas.Timestamp")
-def infer_pandas_timestamp(value):
+def infer_pandas_timestamp(value: pd.Timestamp) -> dt.Timestamp:
     if value.tz is not None:
         return dt.Timestamp(timezone=str(value.tz))
     else:
         return dt.timestamp
 
 
 @infer.register("pandas.Timedelta")
-def infer_interval_pandas(value) -> dt.Interval:
+def infer_interval_pandas(value: pd.Timedelta) -> dt.Interval:
     # pandas Timedelta has more granularity
     units = {
         "D": "d",
@@ -207,7 +254,7 @@ def infer_interval_pandas(value) -> dt.Interval:
 
 @infer.register("numpy.ndarray")
 @infer.register("pandas.Series")
-def infer_numpy_array(value):
+def infer_numpy_array(value: np.ndarray | pd.Series) -> dt.Array:
     from ibis.formats.numpy import NumpyType
     from ibis.formats.pyarrow import PyArrowData
 

ibis/expr/types/arrays.py

@@ -15,6 +15,7 @@
     from collections.abc import Callable, Iterable
 
     import ibis.expr.types as ir
+    from ibis.expr.datatypes.value import InferrableToArray
     from ibis.expr.types.typing import V
 
 import ibis.common.exceptions as com
@@ -1373,6 +1374,36 @@ def means(self) -> ir.FloatingValue:
         """
         return ops.ArrayMean(self).to_expr()
 
+    def __eq__(
+        self, other: InferrableToArray | ArrayValue | Deferred
+    ) -> ir.BooleanValue:
+        return super().__eq__(other)
+
+    def __ne__(
+        self, other: InferrableToArray | ArrayValue | Deferred
+    ) -> ir.BooleanValue:
+        return super().__ne__(other)
+
+    def __ge__(
+        self, other: InferrableToArray | ArrayValue | Deferred
+    ) -> ir.BooleanValue:
+        return super().__ge__(other)
+
+    def __gt__(
+        self, other: InferrableToArray | ArrayValue | Deferred
+    ) -> ir.BooleanValue:
+        return super().__gt__(other)
+
+    def __le__(
+        self, other: InferrableToArray | ArrayValue | Deferred
+    ) -> ir.BooleanValue:
+        return super().__le__(other)
+
+    def __lt__(
+        self, other: InferrableToArray | ArrayValue | Deferred
+    ) -> ir.BooleanValue:
+        return super().__lt__(other)
+
 
 @public
 class ArrayScalar(Scalar, ArrayValue):

ibis/expr/types/generic.py

@@ -31,6 +31,7 @@
 
     import ibis.expr.schema as sch
     import ibis.expr.types as ir
+    from ibis.expr.datatypes.value import Inferrable
     from ibis.formats.pandas import PandasData
     from ibis.formats.pyarrow import PyArrowData
 
@@ -1291,30 +1292,30 @@ def group_concat(
     def __hash__(self) -> int:
         return super().__hash__()
 
-    def __eq__(self, other: Value) -> ir.BooleanValue:
+    def __eq__(self, other: Inferrable | Value | Deferred) -> ir.BooleanValue:
         if _is_null_literal(other):
             return self.isnull()
         elif _is_null_literal(self):
             return other.isnull()
         return _binop(ops.Equals, self, other)
 
-    def __ne__(self, other: Value) -> ir.BooleanValue:
+    def __ne__(self, other: Inferrable | Value | Deferred) -> ir.BooleanValue:
         if _is_null_literal(other):
             return self.notnull()
         elif _is_null_literal(self):
             return other.notnull()
         return _binop(ops.NotEquals, self, other)
 
-    def __ge__(self, other: Value) -> ir.BooleanValue:
+    def __ge__(self, other: Inferrable | Value | Deferred) -> ir.BooleanValue:
         return _binop(ops.GreaterEqual, self, other)
 
-    def __gt__(self, other: Value) -> ir.BooleanValue:
+    def __gt__(self, other: Inferrable | Value | Deferred) -> ir.BooleanValue:
         return _binop(ops.Greater, self, other)
 
-    def __le__(self, other: Value) -> ir.BooleanValue:
+    def __le__(self, other: Inferrable | Value | Deferred) -> ir.BooleanValue:
         return _binop(ops.LessEqual, self, other)
 
-    def __lt__(self, other: Value) -> ir.BooleanValue:
+    def __lt__(self, other: Inferrable | Value | Deferred) -> ir.BooleanValue:
         return _binop(ops.Less, self, other)
 
     def asc(self, *, nulls_first: bool = False) -> Self:
@@ -3045,7 +3046,7 @@ def null(type: dt.DataType | str | None = None, /) -> NullScalar:
 
 @public
 @deferrable
-def literal(value: Any, type: dt.DataType | str | None = None) -> Scalar:
+def literal(value: Inferrable, type: dt.DataType | str | None = None) -> Scalar:
     """Create a scalar expression from a Python value.
 
     ::: {.callout-tip}

ibis/expr/types/numeric.py

@@ -11,14 +11,10 @@
 from ibis.expr.types.generic import Column, Scalar, Value, _binop
 
 if TYPE_CHECKING:
-    from decimal import Decimal
-    from typing import Union
-
     from typing_extensions import Self
 
     import ibis.expr.types as ir
-
-    Number = Union[int, float, Decimal]
+    from ibis.expr.datatypes.value import InferrableToNumeric
 
 
 @public
@@ -170,8 +166,8 @@ def log(self, base: NumericValue | None = None, /) -> NumericValue:
 
     def clip(
         self,
-        lower: Number | NumericValue | ibis.Deferred | None = None,
-        upper: Number | NumericValue | ibis.Deferred | None = None,
+        lower: InferrableToNumeric | NumericValue | ibis.Deferred | None = None,
+        upper: InferrableToNumeric | NumericValue | ibis.Deferred | None = None,
     ) -> Self:
         """Trim values outside of `lower` and `upper` bounds.
 
@@ -537,7 +533,9 @@ def atan(self) -> NumericValue:
         """
         return ops.Atan(self).to_expr()
 
-    def atan2(self, other: Number | NumericValue | ibis.Deferred, /) -> NumericValue:
+    def atan2(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred, /
+    ) -> NumericValue:
         """Compute the two-argument version of arc tangent.
 
         Examples
@@ -642,38 +640,78 @@ def tan(self) -> NumericValue:
         """
         return ops.Tan(self).to_expr()
 
-    def __add__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __eq__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> ir.BooleanValue:
+        return super().__eq__(other)
+
+    def __ne__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> ir.BooleanValue:
+        return super().__ne__(other)
+
+    def __ge__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> ir.BooleanValue:
+        return super().__ge__(other)
+
+    def __gt__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> ir.BooleanValue:
+        return super().__gt__(other)
+
+    def __le__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> ir.BooleanValue:
+        return super().__le__(other)
+
+    def __lt__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> ir.BooleanValue:
+        return super().__lt__(other)
+
+    def __add__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Add `self` with `other`."""
         return _binop(ops.Add, self, other)
 
     add = radd = __radd__ = __add__
 
-    def __sub__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __sub__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Subtract `other` from `self`."""
         return _binop(ops.Subtract, self, other)
 
     sub = __sub__
 
-    def __rsub__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __rsub__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Subtract `self` from `other`."""
         return _binop(ops.Subtract, other, self)
 
     rsub = __rsub__
 
-    def __mul__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __mul__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Multiply `self` and `other`."""
         return _binop(ops.Multiply, self, other)
 
     mul = rmul = __rmul__ = __mul__
 
-    def __truediv__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __truediv__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Divide `self` by `other`."""
         return _binop(ops.Divide, self, other)
 
     div = __div__ = __truediv__
 
     def __rtruediv__(
-        self, other: Number | NumericValue | ibis.Deferred
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
     ) -> NumericValue:
         """Perform `other` / `self`."""
         return _binop(ops.Divide, other, self)
@@ -682,7 +720,7 @@ def __rtruediv__(
 
     def __floordiv__(
         self,
-        other: Number | NumericValue | ibis.Deferred,
+        other: InferrableToNumeric | NumericValue | ibis.Deferred,
     ) -> IntegerValue:
         """Perform `self` // `other`."""
         return _binop(ops.FloorDivide, self, other)
@@ -691,32 +729,40 @@ def __floordiv__(
 
     def __rfloordiv__(
         self,
-        other: Number | NumericValue | ibis.Deferred,
+        other: InferrableToNumeric | NumericValue | ibis.Deferred,
     ) -> IntegerValue:
         """Perform `other` // `self`."""
         return _binop(ops.FloorDivide, other, self)
 
     rfloordiv = __rfloordiv__
 
-    def __pow__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __pow__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Raise `self` to the `other`th power."""
         return _binop(ops.Power, self, other)
 
     pow = __pow__
 
-    def __rpow__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __rpow__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Raise `other` to the `self`th power."""
         return _binop(ops.Power, other, self)
 
     rpow = __rpow__
 
-    def __mod__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __mod__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Compute `self` modulo `other`."""
         return _binop(ops.Modulus, self, other)
 
     mod = __mod__
 
-    def __rmod__(self, other: Number | NumericValue | ibis.Deferred) -> NumericValue:
+    def __rmod__(
+        self, other: InferrableToNumeric | NumericValue | ibis.Deferred
+    ) -> NumericValue:
         """Compute `other` modulo `self`."""
 
         return _binop(ops.Modulus, other, self)