Add truncatable repr for DF groupby groups

Author: benjaminarjun

pandas/core/frame.py

@@ -553,7 +553,7 @@ def _repr_fits_horizontal_(self, ignore_width=False):
         Check if full repr fits in horizontal boundaries imposed by the display
         options width and max_columns.
 
-        In case off non-interactive session, no boundaries apply.
+        In case of non-interactive session, no boundaries apply.
 
         `ignore_width` is here so ipnb+HTML output can behave the way
         users expect. display.max_columns remains in effect.

pandas/core/groupby/groupby.py

@@ -33,7 +33,7 @@ class providing the base-class of operations.
 from pandas.core.base import (
     DataError, GroupByError, PandasObject, SelectionMixin, SpecificationError)
 import pandas.core.common as com
-from pandas.core.config import option_context
+from pandas.core.config import get_option, option_context
 from pandas.core.frame import DataFrame
 from pandas.core.generic import NDFrame
 from pandas.core.groupby import base
@@ -387,7 +387,7 @@ def groups(self):
         Dict {group name -> group labels}.
         """
         self._assure_grouper()
-        return self.grouper.groups
+        return DataFrameGroups(self.grouper.groups)
 
     @property
     def ngroups(self):
@@ -2108,3 +2108,10 @@ def groupby(obj, by, **kwds):
         raise TypeError('invalid type: {}'.format(obj))
 
     return klass(obj, by, **kwds)
+
+
+class DataFrameGroups(dict):
+    def __repr__(self):
+        from pandas.io.formats.printing import _pprint_dict
+        return _pprint_dict(self, max_seq_items=get_option('display.max_rows'),
+                            recurse=False, truncate_at='middle')

pandas/io/formats/printing.py

@@ -95,7 +95,7 @@ def _join_unicode(lines, sep=''):
 def _pprint_seq(seq, _nest_lvl=0, max_seq_items=None, **kwds):
     """
     internal. pprinter for iterables. you should probably use pprint_thing()
-    rather then calling this directly.
+    rather than calling this directly.
 
     bounds length of printed sequence, depending on options
     """
@@ -124,11 +124,13 @@ def _pprint_seq(seq, _nest_lvl=0, max_seq_items=None, **kwds):
     return fmt.format(body=body)
 
 
-def _pprint_dict(seq, _nest_lvl=0, max_seq_items=None, **kwds):
+def _pprint_dict(seq, _nest_lvl=0, max_seq_items=None, recurse=True,
+                 truncate_at='end', **kwds):
     """
     internal. pprinter for iterables. you should probably use pprint_thing()
-    rather then calling this directly.
+    rather than calling this directly.
     """
+
     fmt = u("{{{things}}}")
     pairs = []
 
@@ -139,16 +141,26 @@ def _pprint_dict(seq, _nest_lvl=0, max_seq_items=None, **kwds):
     else:
         nitems = max_seq_items or get_option("max_seq_items") or len(seq)
 
-    for k, v in list(seq.items())[:nitems]:
-        pairs.append(
-            pfmt.format(
-                key=pprint_thing(k, _nest_lvl + 1,
-                                 max_seq_items=max_seq_items, **kwds),
-                val=pprint_thing(v, _nest_lvl + 1,
-                                 max_seq_items=max_seq_items, **kwds)))
+    if recurse:
+        for k, v in list(seq.items())[:nitems]:
+            pairs.append(
+                pfmt.format(
+                    key=pprint_thing(k, _nest_lvl + 1,
+                                     max_seq_items=max_seq_items, **kwds),
+                    val=pprint_thing(v, _nest_lvl + 1,
+                                     max_seq_items=max_seq_items, **kwds)))
+    else:
+        for k, v in list(seq.items())[:nitems]:
+            pairs.append(pfmt.format(key=k, val=v))
 
     if nitems < len(seq):
-        return fmt.format(things=", ".join(pairs) + ", ...")
+        if truncate_at == 'middle':
+            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) + ", ...")
     else:
         return fmt.format(things=", ".join(pairs))
 

pandas/tests/io/formats/test_format.py

@@ -1748,6 +1748,22 @@ def test_period(self):
         assert str(df) == exp
 
 
+class TestDataFrameGroupByFormatting(object):
+    def test_groups_repr_truncates(self):
+        df = pd.DataFrame({
+            'a': [1, 1, 1, 2, 2, 3],
+            'b': [1, 2, 3, 4, 5, 6]
+        })
+
+        with option_context('display.max_rows', 2):
+            x = df.groupby('a').groups
+
+            expected = ("{1: Int64Index([0, 1, 2], dtype='int64'), ... , "
+                        "2: Int64Index([3, 4], dtype='int64')}")
+
+            assert x.__repr__() == expected
+
+
 def gen_series_formatting():
     s1 = pd.Series(['a'] * 100)
     s2 = pd.Series(['ab'] * 100)