overload sort (#2002)

py-polars/polars/internals/frame.py

@@ -1682,12 +1682,43 @@ def replace_at_idx(self, index: int, series: "pli.Series") -> None:
         """
         self._df.replace_at_idx(index, series._s)
 
+    @tp.overload
+    def sort(
+        self,
+        by: Union[str, "pli.Expr", tp.List[str], tp.List["pli.Expr"]],
+        reverse: Union[bool, tp.List[bool]] = ...,
+        *,
+        in_place: Literal[False] = ...,
+    ) -> "DataFrame":
+        ...
+
+    @tp.overload
+    def sort(
+        self,
+        by: Union[str, "pli.Expr", tp.List[str], tp.List["pli.Expr"]],
+        reverse: Union[bool, tp.List[bool]] = ...,
+        *,
+        in_place: Literal[True],
+    ) -> None:
+        ...
+
+    @tp.overload
+    def sort(
+        self,
+        by: Union[str, "pli.Expr", tp.List[str], tp.List["pli.Expr"]],
+        reverse: Union[bool, tp.List[bool]] = ...,
+        *,
+        in_place: bool,
+    ) -> Optional["DataFrame"]:
+        ...
+
     def sort(
         self,
         by: Union[str, "pli.Expr", tp.List[str], tp.List["pli.Expr"]],
         reverse: Union[bool, tp.List[bool]] = False,
+        *,
         in_place: bool = False,
-    ) -> "DataFrame":
+    ) -> Optional["DataFrame"]:
         """
         Sort the DataFrame by column.
 
@@ -1757,7 +1788,7 @@ def sort(
             return df
         if in_place:
             self._df.sort_in_place(by, reverse)
-            return self
+            return None
         else:
             return wrap_df(self._df.sort(by, reverse))
 

py-polars/tests/test_sort.py

@@ -53,3 +53,12 @@ def test_sort_by() -> None:
     # by can also be a single column
     out = df.select([pl.col("a").sort_by("b", reverse=[False])])
     assert out["a"].to_list() == [1, 2, 3, 4, 5]
+
+
+def test_sort_in_place() -> None:
+    df = pl.DataFrame({"a": [1, 3, 2, 4, 5]})
+    ret = df.sort("a", in_place=True)
+    result = df["a"].to_list()
+    expected = [1, 2, 3, 4, 5]
+    assert result == expected
+    assert ret is None