Skip to content

feat(typing): make types more precise for many common Values #11771

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+576 −83
Open

feat(typing): make types more precise for many common Values #11771

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 75 additions & 7 deletions ibis/expr/types/generic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@
class Value(Expr):
"""Base class for a data generating expression having a known type."""

def name(self, name: str, /) -> Value:
def name(self, name: str, /) -> Self:
"""Rename an expression to `name`.

Parameters
Expand Down Expand Up @@ -121,6 +121,10 @@ def type(self) -> dt.DataType:
"""
return self.op().dtype

@overload
def hash(self: Scalar) -> ir.IntegerScalar: ...
@overload
def hash(self: Column) -> ir.IntegerColumn: ...
def hash(self) -> ir.IntegerValue:
"""Compute an integer hash value.

Expand Down Expand Up @@ -297,6 +301,10 @@ def cast(self, target_type: Any, /) -> Value:

return op.to_expr()

@overload
def try_cast(self: Scalar, target_type: Any, /) -> Scalar: ...
@overload
def try_cast(self: Column, target_type: Any, /) -> Column: ...
def try_cast(self, target_type: Any, /) -> Value:
"""Try cast expression to indicated data type.

Expand Down Expand Up @@ -383,6 +391,10 @@ def coalesce(self, /, *args: Value) -> Value:
"""
return ops.Coalesce((self, *args)).to_expr()

@overload
def typeof(self: Scalar) -> ir.StringScalar: ...
@overload
def typeof(self: Column) -> ir.StringColumn: ...
def typeof(self) -> ir.StringValue:
"""Return the string name of the datatype of self.

Expand Down Expand Up @@ -923,6 +935,10 @@ def bind(table):
except com.IbisInputError:
return bind(_)

@overload
def isnull(self: Scalar) -> ir.BooleanScalar: ...
@overload
def isnull(self: Column) -> ir.BooleanColumn: ...
def isnull(self) -> ir.BooleanValue:
"""Whether this expression is `NULL`. Does NOT detect `NaN` and `inf` values.

Expand Down Expand Up @@ -960,6 +976,10 @@ def isnull(self) -> ir.BooleanValue:
"""
return ops.IsNull(self).to_expr()

@overload
def notnull(self: Scalar) -> ir.BooleanScalar: ...
@overload
def notnull(self: Column) -> ir.BooleanColumn: ...
def notnull(self) -> ir.BooleanValue:
"""Return whether this expression is not NULL.

Expand Down Expand Up @@ -1291,29 +1311,77 @@ def group_concat(
def __hash__(self) -> int:
return super().__hash__()

@overload
def __eq__(self: Scalar, other: Scalar) -> ir.BooleanScalar: ...
@overload
def __eq__(self: Scalar, other: Column) -> ir.BooleanColumn: ...
@overload
def __eq__(self: Column, other: Scalar) -> ir.BooleanColumn: ...
@overload
def __eq__(self: Column, other: Column) -> ir.BooleanColumn: ...
def __eq__(self, other: Value) -> ir.BooleanValue:
if _is_null_literal(other):
return self.isnull()
elif _is_null_literal(self):
return other.isnull()
return _binop(ops.Equals, self, other)

@overload
def __ne__(self: Scalar, other: Scalar) -> ir.BooleanScalar: ...
@overload
def __ne__(self: Scalar, other: Column) -> ir.BooleanColumn: ...
@overload
def __ne__(self: Column, other: Scalar) -> ir.BooleanColumn: ...
@overload
def __ne__(self: Column, other: Column) -> ir.BooleanColumn: ...
def __ne__(self, other: Value) -> ir.BooleanValue:
if _is_null_literal(other):
return self.notnull()
elif _is_null_literal(self):
return other.notnull()
return _binop(ops.NotEquals, self, other)

@overload
def __ge__(self: Scalar, other: Scalar) -> ir.BooleanScalar: ...
@overload
def __ge__(self: Scalar, other: Column) -> ir.BooleanColumn: ...
@overload
def __ge__(self: Column, other: Scalar) -> ir.BooleanColumn: ...
@overload
def __ge__(self: Column, other: Column) -> ir.BooleanColumn: ...
def __ge__(self, other: Value) -> ir.BooleanValue:
return _binop(ops.GreaterEqual, self, other)

@overload
def __gt__(self: Scalar, other: Scalar) -> ir.BooleanScalar: ...
@overload
def __gt__(self: Scalar, other: Column) -> ir.BooleanColumn: ...
@overload
def __gt__(self: Column, other: Scalar) -> ir.BooleanColumn: ...
@overload
def __gt__(self: Column, other: Column) -> ir.BooleanColumn: ...
def __gt__(self, other: Value) -> ir.BooleanValue:
return _binop(ops.Greater, self, other)

@overload
def __le__(self: Scalar, other: Scalar) -> ir.BooleanScalar: ...
@overload
def __le__(self: Scalar, other: Column) -> ir.BooleanColumn: ...
@overload
def __le__(self: Column, other: Scalar) -> ir.BooleanColumn: ...
@overload
def __le__(self: Column, other: Column) -> ir.BooleanColumn: ...
def __le__(self, other: Value) -> ir.BooleanValue:
return _binop(ops.LessEqual, self, other)

@overload
def __lt__(self: Scalar, other: Scalar) -> ir.BooleanScalar: ...
@overload
def __lt__(self: Scalar, other: Column) -> ir.BooleanColumn: ...
@overload
def __lt__(self: Column, other: Scalar) -> ir.BooleanColumn: ...
@overload
def __lt__(self: Column, other: Column) -> ir.BooleanColumn: ...
def __lt__(self, other: Value) -> ir.BooleanValue:
return _binop(ops.Less, self, other)

Expand Down Expand Up @@ -1484,7 +1552,7 @@ def __polars_result__(self, df: pl.DataFrame) -> Any:

return PolarsData.convert_scalar(df, self.type())

def as_scalar(self):
def as_scalar(self) -> Self:
"""Inform ibis that the expression should be treated as a scalar.

If the expression is a literal, it will be returned as is. If it depends
Expand Down Expand Up @@ -2461,7 +2529,7 @@ def first(
where: ir.BooleanValue | None = None,
order_by: Any = None,
include_null: bool = False,
) -> Value:
) -> Scalar:
"""Return the first value of a column.

Parameters
Expand Down Expand Up @@ -2515,7 +2583,7 @@ def last(
where: ir.BooleanValue | None = None,
order_by: Any = None,
include_null: bool = False,
) -> Value:
) -> Scalar:
"""Return the last value of a column.

Parameters
Expand Down Expand Up @@ -2627,7 +2695,7 @@ def dense_rank(self) -> ir.IntegerColumn:
"""
return ibis.dense_rank().over(order_by=self)

def percent_rank(self) -> Column:
def percent_rank(self) -> ir.IntegerColumn:
"""Return the relative rank of the values in the column.

Examples
Expand All @@ -2651,7 +2719,7 @@ def percent_rank(self) -> Column:
"""
return ibis.percent_rank().over(order_by=self)

def cume_dist(self) -> Column:
def cume_dist(self) -> ir.FloatingColumn:
"""Return the cumulative distribution over a window.

Examples
Expand Down Expand Up @@ -2995,7 +3063,7 @@ class NullColumn(Column, NullValue):

@public
@deferrable
def null(type: dt.DataType | str | None = None, /) -> Value:
def null(type: dt.DataType | str | None = None, /) -> Scalar:
"""Create a NULL scalar.

`NULL`s with an unspecified type are castable and comparable to values,
Expand Down
40 changes: 30 additions & 10 deletions ibis/expr/types/logical.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, overload

from public import public

Expand All @@ -11,11 +11,19 @@
from ibis.expr.types.numeric import NumericColumn, NumericScalar, NumericValue

if TYPE_CHECKING:
from typing_extensions import Self

import ibis.expr.types as ir


@public
class BooleanValue(NumericValue):
@overload
def ifelse(
self: ir.BooleanScalar, true_expr: ir.Scalar, false_expr: ir.Scalar, /
) -> ir.Scalar: ...
@overload
def ifelse(self, true_expr: ir.Value, false_expr: ir.Value, /) -> ir.Column: ...
def ifelse(self, true_expr: ir.Value, false_expr: ir.Value, /) -> ir.Value:
"""Construct a ternary conditional expression.

Expand Down Expand Up @@ -53,7 +61,11 @@ def ifelse(self, true_expr: ir.Value, false_expr: ir.Value, /) -> ir.Value:
# must be used.
return ops.IfElse(self, true_expr, false_expr).to_expr()

def __and__(self, other: BooleanValue) -> BooleanValue:
@overload
def __and__(self: BooleanScalar, other: bool | BooleanScalar) -> BooleanScalar: ...
@overload
def __and__(self, other: bool | BooleanValue) -> BooleanColumn: ...
def __and__(self, other: bool | BooleanValue) -> BooleanValue:
"""Construct a binary AND conditional expression with `self` and `other`.

Parameters
Expand Down Expand Up @@ -101,7 +113,11 @@ def __and__(self, other: BooleanValue) -> BooleanValue:

__rand__ = __and__

def __or__(self, other: BooleanValue) -> BooleanValue:
@overload
def __or__(self: BooleanScalar, other: bool | BooleanScalar) -> BooleanScalar: ...
@overload
def __or__(self, other: bool | BooleanValue) -> BooleanColumn: ...
def __or__(self, other: bool | BooleanValue) -> BooleanValue:
"""Construct a binary OR conditional expression with `self` and `other`.

Parameters
Expand Down Expand Up @@ -137,7 +153,11 @@ def __or__(self, other: BooleanValue) -> BooleanValue:

__ror__ = __or__

def __xor__(self, other: BooleanValue) -> BooleanValue:
@overload
def __xor__(self: BooleanScalar, other: bool | BooleanScalar) -> BooleanScalar: ...
@overload
def __xor__(self, other: bool | BooleanValue) -> BooleanColumn: ...
def __xor__(self, other: bool | BooleanValue) -> BooleanValue:
"""Construct a binary XOR conditional expression with `self` and `other`.

Parameters
Expand Down Expand Up @@ -198,7 +218,7 @@ def __xor__(self, other: BooleanValue) -> BooleanValue:

__rxor__ = __xor__

def __invert__(self) -> BooleanValue:
def __invert__(self) -> Self:
"""Construct a unary NOT conditional expression with `self`.

Parameters
Expand Down Expand Up @@ -230,7 +250,7 @@ def __invert__(self) -> BooleanValue:
"""
return ops.Not(self).to_expr()

def negate(self) -> BooleanValue:
def negate(self) -> Self:
"""DEPRECATED."""
util.warn_deprecated(
"`-bool_val`/`bool_val.negate()`",
Expand All @@ -247,7 +267,7 @@ class BooleanScalar(NumericScalar, BooleanValue):

@public
class BooleanColumn(NumericColumn, BooleanValue):
def any(self, *, where: BooleanValue | None = None) -> BooleanValue:
def any(self, *, where: bool | BooleanValue | None = None) -> BooleanScalar:
"""Return whether at least one element is `True`.

If the expression does not reference any foreign tables, the result
Expand Down Expand Up @@ -339,7 +359,7 @@ def resolve_exists_subquery(outer):

return op.to_expr()

def notany(self, *, where: BooleanValue | None = None) -> BooleanValue:
def notany(self, *, where: bool | BooleanValue | None = None) -> BooleanScalar:
"""Return whether no elements are `True`.

Parameters
Expand Down Expand Up @@ -373,7 +393,7 @@ def notany(self, *, where: BooleanValue | None = None) -> BooleanValue:
"""
return ~self.any(where=where)

def all(self, *, where: BooleanValue | None = None) -> BooleanScalar:
def all(self, *, where: bool | BooleanValue | None = None) -> BooleanScalar:
"""Return whether all elements are `True`.

Parameters
Expand Down Expand Up @@ -410,7 +430,7 @@ def all(self, *, where: BooleanValue | None = None) -> BooleanScalar:
"""
return ops.All(self, where=self._bind_to_parent_table(where)).to_expr()

def notall(self, *, where: BooleanValue | None = None) -> BooleanScalar:
def notall(self, *, where: bool | BooleanValue | None = None) -> BooleanScalar:
"""Return whether not all elements are `True`.

Parameters
Expand Down
Loading