diff --git a/ibis/expr/datatypes/value.py b/ibis/expr/datatypes/value.py index 30ad6d26411e..ee426b76276d 100644 --- a/ibis/expr/datatypes/value.py +++ b/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,14 +216,14 @@ 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: @@ -184,7 +231,7 @@ def infer_pandas_timestamp(value): @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 diff --git a/ibis/expr/types/arrays.py b/ibis/expr/types/arrays.py index addf296314ee..b2c5f9faf346 100644 --- a/ibis/expr/types/arrays.py +++ b/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): diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py index 6a6141ade383..fca28456f6d8 100644 --- a/ibis/expr/types/generic.py +++ b/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} diff --git a/ibis/expr/types/numeric.py b/ibis/expr/types/numeric.py index 19f324c0ac12..03cf612f0f7f 100644 --- a/ibis/expr/types/numeric.py +++ b/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) diff --git a/ibis/expr/types/strings.py b/ibis/expr/types/strings.py index 3891e2b89b50..e72c4ed07e2d 100644 --- a/ibis/expr/types/strings.py +++ b/ibis/expr/types/strings.py @@ -17,6 +17,7 @@ import ibis.expr.types as ir from ibis.common.deferred import Deferred + from ibis.expr.datatypes import InferrableToString @public @@ -1605,7 +1606,39 @@ def concat( """ return ops.StringConcat((self, other, *args)).to_expr() - def __add__(self, other: str | StringValue | Deferred) -> StringValue: + def __eq__( + self, other: InferrableToString | StringValue | Deferred + ) -> ir.BooleanValue: + return super().__eq__(other) + + def __ne__( + self, other: InferrableToString | StringValue | Deferred + ) -> ir.BooleanValue: + return super().__ne__(other) + + def __ge__( + self, other: InferrableToString | StringValue | Deferred + ) -> ir.BooleanValue: + return super().__ge__(other) + + def __gt__( + self, other: InferrableToString | StringValue | Deferred + ) -> ir.BooleanValue: + return super().__gt__(other) + + def __le__( + self, other: InferrableToString | StringValue | Deferred + ) -> ir.BooleanValue: + return super().__le__(other) + + def __lt__( + self, other: InferrableToString | StringValue | Deferred + ) -> ir.BooleanValue: + return super().__lt__(other) + + def __add__( + self, other: InferrableToString | StringValue | Deferred + ) -> StringValue: """Concatenate strings. Parameters diff --git a/ibis/expr/types/uuid.py b/ibis/expr/types/uuid.py index 772e74da5186..75c7d38b726e 100644 --- a/ibis/expr/types/uuid.py +++ b/ibis/expr/types/uuid.py @@ -1,13 +1,78 @@ from __future__ import annotations +from typing import TYPE_CHECKING + from public import public from ibis.expr.types.generic import Column, Scalar, Value +if TYPE_CHECKING: + from ibis.common.deferred import Deferred + from ibis.expr import types as ir + from ibis.expr.datatypes.value import InferrableToString, InferrableToUUID + @public class UUIDValue(Value): - pass + def __eq__( + self, + other: InferrableToUUID + | UUIDValue + | InferrableToString + | ir.StringValue + | Deferred, + ) -> ir.BooleanValue: + return super().__eq__(other) + + def __ne__( + self, + other: InferrableToUUID + | UUIDValue + | InferrableToString + | ir.StringValue + | Deferred, + ) -> ir.BooleanValue: + return super().__ne__(other) + + def __ge__( + self, + other: InferrableToUUID + | UUIDValue + | InferrableToString + | ir.StringValue + | Deferred, + ) -> ir.BooleanValue: + return super().__ge__(other) + + def __gt__( + self, + other: InferrableToUUID + | UUIDValue + | InferrableToString + | ir.StringValue + | Deferred, + ) -> ir.BooleanValue: + return super().__gt__(other) + + def __le__( + self, + other: InferrableToUUID + | UUIDValue + | InferrableToString + | ir.StringValue + | Deferred, + ) -> ir.BooleanValue: + return super().__le__(other) + + def __lt__( + self, + other: InferrableToUUID + | UUIDValue + | InferrableToString + | ir.StringValue + | Deferred, + ) -> ir.BooleanValue: + return super().__lt__(other) @public