Skip to content

Commit

Permalink
fix(python): Change method op type hints to match dunder op type hints
Browse files Browse the repository at this point in the history
Type hints for operators were refined in #13635 [0], but the refinements
were generally not propagated to the method equivalents (pow() being the
sole exception). This can result in equivalent expressions being treated
differently by type checkers, especially if stricter settings are used
that discourage/disallow the use of Any.

This commit changes the type hints for the method equivalents to match
those for the operators. This should ensure resulting types are
consistent regardless of whether one uses infix operators or their
method equivalents.

This also has the side benefit of making the documentation a bit more
consistent as the type hint more closely matches the stated type of the
argument in the docstring.

Not entirely sure whether the type: ignore hints are necessary or
whether there's some alternative formulation which doesn't require them.
I believe the override should be sound in this instance since the
expr.or_(expr, ...) should be equivalent to expr | expr | ..., and the
type signature for or_() was changed to be the same as for __bitor__
(and a similar argument applies to and_()/__bitand__).

[0]: 1c417e4
  • Loading branch information
Alex Wang committed Oct 22, 2024
1 parent e8cfa44 commit 604988d
Showing 1 changed file with 19 additions and 19 deletions.
38 changes: 19 additions & 19 deletions py-polars/polars/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -4965,7 +4965,7 @@ def limit(self, n: int | Expr = 10) -> Expr:
"""
return self.head(n)

def and_(self, *others: Any) -> Expr:
def and_(self, *others: IntoExprColumn | int | bool) -> Expr:
"""
Method equivalent of bitwise "and" operator `expr & other & ...`.
Expand Down Expand Up @@ -5006,9 +5006,9 @@ def and_(self, *others: Any) -> Expr:
│ false │
└───────┘
"""
return reduce(operator.and_, (self, *others))
return reduce(operator.and_, (self, *others)) # type: ignore[return-value]

def or_(self, *others: Any) -> Expr:
def or_(self, *others: IntoExprColumn | int | bool) -> Expr:
"""
Method equivalent of bitwise "or" operator `expr | other | ...`.
Expand Down Expand Up @@ -5048,9 +5048,9 @@ def or_(self, *others: Any) -> Expr:
│ false │
└───────┘
"""
return reduce(operator.or_, (self,) + others)
return reduce(operator.or_, (self,) + others) # type: ignore[return-value]

def eq(self, other: Any) -> Expr:
def eq(self, other: IntoExpr) -> Expr:
"""
Method equivalent of equality operator `expr == other`.
Expand Down Expand Up @@ -5084,7 +5084,7 @@ def eq(self, other: Any) -> Expr:
"""
return self.__eq__(other)

def eq_missing(self, other: Any) -> Expr:
def eq_missing(self, other: IntoExpr) -> Expr:
"""
Method equivalent of equality operator `expr == other` where `None == None`.
Expand Down Expand Up @@ -5124,7 +5124,7 @@ def eq_missing(self, other: Any) -> Expr:
other = parse_into_expression(other, str_as_lit=True)
return self._from_pyexpr(self._pyexpr.eq_missing(other))

def ge(self, other: Any) -> Expr:
def ge(self, other: IntoExpr) -> Expr:
"""
Method equivalent of "greater than or equal" operator `expr >= other`.
Expand Down Expand Up @@ -5158,7 +5158,7 @@ def ge(self, other: Any) -> Expr:
"""
return self.__ge__(other)

def gt(self, other: Any) -> Expr:
def gt(self, other: IntoExpr) -> Expr:
"""
Method equivalent of "greater than" operator `expr > other`.
Expand Down Expand Up @@ -5192,7 +5192,7 @@ def gt(self, other: Any) -> Expr:
"""
return self.__gt__(other)

def le(self, other: Any) -> Expr:
def le(self, other: IntoExpr) -> Expr:
"""
Method equivalent of "less than or equal" operator `expr <= other`.
Expand Down Expand Up @@ -5226,7 +5226,7 @@ def le(self, other: Any) -> Expr:
"""
return self.__le__(other)

def lt(self, other: Any) -> Expr:
def lt(self, other: IntoExpr) -> Expr:
"""
Method equivalent of "less than" operator `expr < other`.
Expand Down Expand Up @@ -5260,7 +5260,7 @@ def lt(self, other: Any) -> Expr:
"""
return self.__lt__(other)

def ne(self, other: Any) -> Expr:
def ne(self, other: IntoExpr) -> Expr:
"""
Method equivalent of inequality operator `expr != other`.
Expand Down Expand Up @@ -5294,7 +5294,7 @@ def ne(self, other: Any) -> Expr:
"""
return self.__ne__(other)

def ne_missing(self, other: Any) -> Expr:
def ne_missing(self, other: IntoExpr) -> Expr:
"""
Method equivalent of equality operator `expr != other` where `None == None`.
Expand Down Expand Up @@ -5334,7 +5334,7 @@ def ne_missing(self, other: Any) -> Expr:
other = parse_into_expression(other, str_as_lit=True)
return self._from_pyexpr(self._pyexpr.neq_missing(other))

def add(self, other: Any) -> Expr:
def add(self, other: IntoExpr) -> Expr:
"""
Method equivalent of addition operator `expr + other`.
Expand Down Expand Up @@ -5380,7 +5380,7 @@ def add(self, other: Any) -> Expr:
"""
return self.__add__(other)

def floordiv(self, other: Any) -> Expr:
def floordiv(self, other: IntoExpr) -> Expr:
"""
Method equivalent of integer division operator `expr // other`.
Expand Down Expand Up @@ -5467,7 +5467,7 @@ def floordiv(self, other: Any) -> Expr:
"""
return self.__floordiv__(other)

def mod(self, other: Any) -> Expr:
def mod(self, other: IntoExpr) -> Expr:
"""
Method equivalent of modulus operator `expr % other`.
Expand Down Expand Up @@ -5495,7 +5495,7 @@ def mod(self, other: Any) -> Expr:
"""
return self.__mod__(other)

def mul(self, other: Any) -> Expr:
def mul(self, other: IntoExpr) -> Expr:
"""
Method equivalent of multiplication operator `expr * other`.
Expand Down Expand Up @@ -5526,7 +5526,7 @@ def mul(self, other: Any) -> Expr:
"""
return self.__mul__(other)

def sub(self, other: Any) -> Expr:
def sub(self, other: IntoExpr) -> Expr:
"""
Method equivalent of subtraction operator `expr - other`.
Expand Down Expand Up @@ -5579,7 +5579,7 @@ def neg(self) -> Expr:
"""
return self.__neg__()

def truediv(self, other: Any) -> Expr:
def truediv(self, other: IntoExpr) -> Expr:
"""
Method equivalent of float division operator `expr / other`.
Expand Down Expand Up @@ -5676,7 +5676,7 @@ def pow(self, exponent: IntoExprColumn | int | float) -> Expr:
"""
return self.__pow__(exponent)

def xor(self, other: Any) -> Expr:
def xor(self, other: IntoExprColumn | int | bool) -> Expr:
"""
Method equivalent of bitwise exclusive-or operator `expr ^ other`.
Expand Down

0 comments on commit 604988d

Please sign in to comment.