From 604988ddafa44450e411f228abae1c0349dc4046 Mon Sep 17 00:00:00 2001
From: Alex Wang <ts826848@gmail.com>
Date: Wed, 28 Aug 2024 14:45:05 -0400
Subject: [PATCH] fix(python): Change method op type hints to match dunder op
 type hints

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]: https://github.com/pola-rs/polars/pull/13635/commits/1c417e44cd7297d89a6931c50b3df565d3ee04e4
---
 py-polars/polars/expr/expr.py | 38 +++++++++++++++++------------------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/py-polars/polars/expr/expr.py b/py-polars/polars/expr/expr.py
index d5d40217f72d..ec308fa9c28c 100644
--- a/py-polars/polars/expr/expr.py
+++ b/py-polars/polars/expr/expr.py
@@ -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 & ...`.
 
@@ -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 | ...`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.
 
@@ -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`.