Disable getitem for GroupBy (#4330)

pola-rs · Aug 9, 2022 · d74f306 · d74f306
1 parent 43bee70
commit d74f306
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 27 deletions.
diff --git a/py-polars/docs/source/reference/dataframe.rst b/py-polars/docs/source/reference/dataframe.rst
@@ -160,7 +160,6 @@ This namespace comes available by calling `DataFrame.groupby(..)`.
     GroupBy.apply
     GroupBy.count
     GroupBy.first
-    GroupBy.groups
     GroupBy.head
     GroupBy.last
     GroupBy.max

diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -6049,11 +6049,12 @@ def __init__(
         self.by = by
         self.maintain_order = maintain_order
 
-    def __getitem__(self, item: Any) -> GBSelection[DF]:
-        print(
-            "accessing GroupBy by index is deprecated, consider using the `.agg` method"
-        )
-        return self._select(item)
+    def __iter__(self) -> Iterable[Any]:
+        groups_df = self._groups()
+        groups = groups_df["groups"]
+        df = self._dataframe_class._from_pydf(self._df)
+        for i in range(groups_df.height):
+            yield df[groups[i]]
 
     def _select(self, columns: str | list[str]) -> GBSelection[DF]:  # pragma: no cover
         """
@@ -6079,14 +6080,16 @@ def _select(self, columns: str | list[str]) -> GBSelection[DF]:  # pragma: no co
             dataframe_class=self._dataframe_class,
         )
 
-    def __iter__(self) -> Iterable[Any]:
-        groups_df = self.groups()
-        groups = groups_df["groups"]
-        df = self._dataframe_class._from_pydf(self._df)
-        for i in range(groups_df.height):
-            yield df[groups[i]]
+    def _select_all(self) -> GBSelection[DF]:
+        """Select all columns for aggregation."""
+        return GBSelection(
+            self._df,
+            self.by,
+            None,
+            dataframe_class=self._dataframe_class,
+        )
 
-    def groups(self) -> DF:  # pragma: no cover
+    def _groups(self) -> DF:  # pragma: no cover
         """
         Return a `DataFrame` with:
 
@@ -6104,7 +6107,7 @@ def groups(self) -> DF:  # pragma: no cover
         ...     }
         ... )
 
-        >>> df.groupby("d").groups().sort(by="d")
+        >>> df.groupby("d")._groups().sort(by="d")
         shape: (3, 2)
         ┌────────┬───────────┐
         │ d      ┆ groups    │
@@ -6382,15 +6385,6 @@ def tail(self, n: int = 5) -> DF:
         )
         return self._dataframe_class._from_pydf(df._df)
 
-    def _select_all(self) -> GBSelection[DF]:
-        """Select all columns for aggregation."""
-        return GBSelection(
-            self._df,
-            self.by,
-            None,
-            dataframe_class=self._dataframe_class,
-        )
-
     def pivot(
         self, pivot_column: str | list[str], values_column: str | list[str]
     ) -> PivotOps[DF]:

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -343,10 +343,6 @@ def test_groupby() -> None:
 
     assert df.groupby("a").apply(lambda df: df[["c"]].sum()).sort("c")["c"][0] == 1
 
-    with pytest.deprecated_call():
-        df_groups = df.groupby("a").groups().sort("a")
-        assert df_groups["a"].series_equal(pl.Series("a", ["a", "b", "c"]))
-
     with pytest.deprecated_call():
         # TODO: find a way to avoid indexing into GroupBy
         for subdf in df.groupby("a"):  # type: ignore[attr-defined]