Skip to content

Commit

Permalink
feat(python): typing overloads for Series operator methods `ge, gt, .…
Browse files Browse the repository at this point in the history
…..` (pola-rs#13167)

Co-authored-by: s-banach <john@hopfensperger.family>
  • Loading branch information
2 people authored and bushuyev committed Jan 11, 2024
1 parent 55a3ad1 commit 7018106
Showing 1 changed file with 48 additions and 0 deletions.
48 changes: 48 additions & 0 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -782,14 +782,38 @@ def __le__(self, other: Any) -> Series | Expr:
return F.lit(self).__le__(other)
return self._comp(other, "lt_eq")

@overload
def le(self, other: Expr) -> Expr: # type: ignore[overload-overlap]
...

@overload
def le(self, other: Any) -> Series:
...

def le(self, other: Any) -> Self | Expr:
"""Method equivalent of operator expression `series <= other`."""
return self.__le__(other)

@overload
def lt(self, other: Expr) -> Expr: # type: ignore[overload-overlap]
...

@overload
def lt(self, other: Any) -> Series:
...

def lt(self, other: Any) -> Self | Expr:
"""Method equivalent of operator expression `series < other`."""
return self.__lt__(other)

@overload
def eq(self, other: Expr) -> Expr: # type: ignore[overload-overlap]
...

@overload
def eq(self, other: Any) -> Series:
...

def eq(self, other: Any) -> Self | Expr:
"""Method equivalent of operator expression `series == other`."""
return self.__eq__(other)
Expand Down Expand Up @@ -840,6 +864,14 @@ def eq_missing(self, other: Any) -> Self | Expr:
]
"""

@overload
def ne(self, other: Expr) -> Expr: # type: ignore[overload-overlap]
...

@overload
def ne(self, other: Any) -> Series:
...

def ne(self, other: Any) -> Self | Expr:
"""Method equivalent of operator expression `series != other`."""
return self.__ne__(other)
Expand Down Expand Up @@ -890,10 +922,26 @@ def ne_missing(self, other: Any) -> Self | Expr:
]
"""

@overload
def ge(self, other: Expr) -> Expr: # type: ignore[overload-overlap]
...

@overload
def ge(self, other: Any) -> Series:
...

def ge(self, other: Any) -> Self | Expr:
"""Method equivalent of operator expression `series >= other`."""
return self.__ge__(other)

@overload
def gt(self, other: Expr) -> Expr: # type: ignore[overload-overlap]
...

@overload
def gt(self, other: Any) -> Series:
...

def gt(self, other: Any) -> Self | Expr:
"""Method equivalent of operator expression `series > other`."""
return self.__gt__(other)
Expand Down

0 comments on commit 7018106

Please sign in to comment.