Move truncated dict repr to Index.groupby()

Author: benjaminarjun

pandas/core/groupby/base.py

@@ -10,7 +10,6 @@
 
 from pandas.core.dtypes.common import is_list_like, is_scalar
 
-
 class GroupByMixin(object):
     """
     Provide the groupby facilities to the mixed object.

pandas/core/groupby/groupby.py

@@ -390,7 +390,7 @@ def groups(self):
         Dict {group name -> group labels}.
         """
         self._assure_grouper()
-        return DataFrameGroups(self.grouper.groups)
+        return self.grouper.groups
 
     @property
     def ngroups(self):
@@ -2204,25 +2204,3 @@ def groupby(obj, by, **kwds):
         raise TypeError('invalid type: {}'.format(obj))
 
     return klass(obj, by, **kwds)
-
-
-class DataFrameGroups(dict):
-    def __repr__(self):
-        from pandas.compat import u
-
-        nitems = get_option('display.max_rows') or len(self)
-
-        fmt = u("{{{things}}}")
-        pfmt = u("{key}: {val}")
-
-        pairs = []
-        for k, v in list(self.items()):
-            pairs.append(pfmt.format(key=k, val=v))
-
-        if nitems < len(self):
-            start_cnt, end_cnt = nitems - int(nitems / 2), int(nitems / 2)
-            return fmt.format(things=", ".join(pairs[:start_cnt]) +
-                                     ", ... , " +
-                                     ", ".join(pairs[-end_cnt:]))
-        else:
-            return fmt.format(things=", ".join(pairs))

pandas/core/groupby/grouper.py

@@ -236,7 +236,6 @@ class Grouping(object):
 
     def __init__(self, index, grouper=None, obj=None, name=None, level=None,
                  sort=True, observed=False, in_axis=False):
-
         self.name = name
         self.level = level
         self.grouper = _convert_grouper(index, grouper)

pandas/core/groupby/ops.py

@@ -235,12 +235,13 @@ def size(self):
     @cache_readonly
     def groups(self):
         """ dict {group name -> group labels} """
+
         if len(self.groupings) == 1:
             return self.groupings[0].groups
         else:
             to_groupby = lzip(*(ping.grouper for ping in self.groupings))
             to_groupby = Index(to_groupby)
-            return self.axis.groupby(to_groupby)
+            return BaseGrouperGroups(self.axis.groupby(to_groupby))
 
     @cache_readonly
     def is_monotonic(self):
@@ -361,7 +362,7 @@ def get_group_levels(self):
 
     def _is_builtin_func(self, arg):
         """
-        if we define an builtin function for this argument, return it,
+        if we define a builtin function for this argument, return it,
         otherwise return the arg
         """
         return SelectionMixin._builtin_table.get(arg, arg)

pandas/core/indexes/base.py

@@ -38,6 +38,7 @@
 from pandas.core.arrays import ExtensionArray
 from pandas.core.base import IndexOpsMixin, PandasObject
 import pandas.core.common as com
+from pandas.core.config import get_option
 from pandas.core.indexes.frozen import FrozenList
 import pandas.core.missing as missing
 from pandas.core.ops import get_op_result_name, make_invalid_op
@@ -4493,7 +4494,7 @@ def groupby(self, values):
         # map to the label
         result = {k: self.take(v) for k, v in compat.iteritems(result)}
 
-        return result
+        return BaseGroupbyGroups(result)
 
     def map(self, mapper, na_action=None):
         """
@@ -5290,6 +5291,29 @@ def _add_logical_methods_disabled(cls):
 Index._add_comparison_methods()
 
 
+class IndexGroupbyGroups(dict):
+    def __repr__(self):
+        from pandas.compat import u
+
+        nitems = get_option('display.max_rows') or len(self)
+
+        fmt = u("{{{things}}}")
+        pfmt = u("{key}: {val}")
+
+        pairs = []
+        for k, v in list(self.items()):
+            pairs.append(pfmt.format(key=k, val=v))
+
+        if nitems < len(self):
+            print("Truncating repr")
+            start_cnt, end_cnt = nitems - int(nitems / 2), int(nitems / 2)
+            return fmt.format(things=", ".join(pairs[:start_cnt]) +
+                                     ", ... , " +
+                                     ", ".join(pairs[-end_cnt:]))
+        else:
+            return fmt.format(things=", ".join(pairs))
+
+
 def ensure_index_from_sequences(sequences, names=None):
     """
     Construct an index from sequences of data.

pandas/tests/groupby/test_grouping.py

@@ -667,21 +667,21 @@ def test_gb_key_len_equal_axis_len(self):
 
 class TestIteration():
 
-    # def test_groups(self, df):
-    #     grouped = df.groupby(['A'])
-    #     groups = grouped.groups
-    #     assert groups is grouped.groups  # caching works
-    #
-    #     for k, v in compat.iteritems(grouped.groups):
-    #         assert (df.loc[v]['A'] == k).all()
-    #
-    #     grouped = df.groupby(['A', 'B'])
-    #     groups = grouped.groups
-    #     assert groups is grouped.groups  # caching works
-    #
-    #     for k, v in compat.iteritems(grouped.groups):
-    #         assert (df.loc[v]['A'] == k[0]).all()
-    #         assert (df.loc[v]['B'] == k[1]).all()
+    def test_groups(self, df):
+        grouped = df.groupby(['A'])
+        groups = grouped.groups
+        assert groups is grouped.groups  # caching works
+
+        for k, v in compat.iteritems(grouped.groups):
+            assert (df.loc[v]['A'] == k).all()
+
+        grouped = df.groupby(['A', 'B'])
+        groups = grouped.groups
+        assert groups is grouped.groups  # caching works
+
+        for k, v in compat.iteritems(grouped.groups):
+            assert (df.loc[v]['A'] == k[0]).all()
+            assert (df.loc[v]['B'] == k[1]).all()
 
     def test_grouping_is_iterable(self, tsframe):
         # this code path isn't used anywhere else