rewrite tests, typing

Author: Marco Gorelli

pandas/core/indexes/base.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 import operator
 from textwrap import dedent
-from typing import Any, Dict, FrozenSet, Hashable, Optional, Union
+from typing import Any, FrozenSet, Hashable, Optional, Union
 import warnings
 
 import numpy as np
@@ -4766,7 +4766,7 @@ def _maybe_promote(self, other):
                 return self.astype("object"), other.astype("object")
         return self, other
 
-    def groupby(self, values) -> Dict[Hashable, np.ndarray]:
+    def groupby(self, values) -> PrettyDict:
         """
         Group the index labels by a given array of values.
 

pandas/io/formats/printing.py

@@ -6,12 +6,14 @@
 from typing import (
     Any,
     Callable,
+    Dict,
     Iterable,
     List,
     Mapping,
     Optional,
     Sequence,
     Tuple,
+    TypeVar,
     Union,
 )
 
@@ -20,6 +22,8 @@
 from pandas.core.dtypes.inference import is_sequence
 
 EscapeChars = Union[Mapping[str, str], Iterable[str]]
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
 
 
 def adjoin(space: int, *lists: List[str], **kwargs) -> str:
@@ -530,8 +534,8 @@ def format_object_attrs(
     return attrs
 
 
-class PrettyDict(dict):
+class PrettyDict(Dict[_KT, _VT]):
     """Dict extension to support abbreviated __repr__"""
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return pprint_thing(self, max_seq_items=get_option("display.max_rows"))

pandas/tests/groupby/test_groupby.py

@@ -2041,12 +2041,12 @@ def test_groupby_list_level():
 
 def test_groups_repr_truncates():
     # GH 1135
-    df = pd.DataFrame({"a": [1, 1, 1, 2, 2, 3], "b": [1, 2, 3, 4, 5, 6]})
+    df = pd.DataFrame(np.random.randn(61, 1))
+    df["a"] = df.index
 
-    with pd.option_context("display.max_rows", 2):
-        x = df.groupby("a").groups
-        assert x.__repr__().endswith("...}")
+    result = df.groupby("a").groups.__repr__()
+    expected = str({i: [i] for i in range(60)})[:-1] + ", ...}"
+    assert result == expected
 
-    with pd.option_context("display.max_rows", 5):
-        x = df.groupby(np.array(df.a)).groups
-        assert not x.__repr__().endswith("...}")
+    result = df.groupby(np.array(df.a)).groups.__repr__()
+    assert result == expected