Add typing.Counter and typing.ChainMap (#366)

JelleZijlstra · ilevkivskyi · commit d6390f31719a · 2017-02-10T20:18:10.000+01:00
* Add typing.Counter

Pretty much copied from deque.

* also add ChainMap

Only to the Python 3 version because ChainMap was added in 3.3.

* fix for 3.2

Apparently typing still supports 3.2

* fix flake8

* fix base class of Counter

* move counter below Dict

* clarify error message, add test

* allow instantiation

* fix reliance on type cache
diff --git a/python2/test_typing.py b/python2/test_typing.py
@@ -1284,6 +1284,9 @@ def test_deque(self):
         class MyDeque(typing.Deque[int]): pass
         self.assertIsInstance(MyDeque(), collections.deque)
 
+    def test_counter(self):
+        self.assertIsSubclass(collections.Counter, typing.Counter)
+
     def test_set(self):
         self.assertIsSubclass(set, typing.Set)
         self.assertNotIsSubclass(frozenset, typing.Set)
@@ -1358,6 +1361,23 @@ def test_deque_instantiation(self):
         class D(typing.Deque[T]): pass
         self.assertIs(type(D[int]()), D)
 
+    def test_counter_instantiation(self):
+        self.assertIs(type(typing.Counter()), collections.Counter)
+        self.assertIs(type(typing.Counter[T]()), collections.Counter)
+        self.assertIs(type(typing.Counter[int]()), collections.Counter)
+        class C(typing.Counter[T]): pass
+        self.assertIs(type(C[int]()), C)
+
+    def test_counter_subclass_instantiation(self):
+
+        class MyCounter(typing.Counter[int]):
+            pass
+
+        d = MyCounter()
+        self.assertIsInstance(d, MyCounter)
+        self.assertIsInstance(d, typing.Counter)
+        self.assertIsInstance(d, collections.Counter)
+
     def test_no_set_instantiation(self):
         with self.assertRaises(TypeError):
             typing.Set()
diff --git a/python2/typing.py b/python2/typing.py
@@ -53,6 +53,7 @@
     'SupportsInt',
 
     # Concrete collection types.
+    'Counter',
     'Deque',
     'Dict',
     'DefaultDict',
@@ -1785,6 +1786,16 @@ def __new__(cls, *args, **kwds):
         return _generic_new(collections.defaultdict, cls, *args, **kwds)
 
 
+class Counter(collections.Counter, Dict[T, int]):
+    __slots__ = ()
+    __extra__ = collections.Counter
+
+    def __new__(cls, *args, **kwds):
+        if _geqv(cls, Counter):
+            return collections.Counter(*args, **kwds)
+        return _generic_new(collections.Counter, cls, *args, **kwds)
+
+
 # Determine what base class to use for Generator.
 if hasattr(collections_abc, 'Generator'):
     # Sufficiently recent versions of 3.5 have a Generator ABC.
diff --git a/src/test_typing.py b/src/test_typing.py
@@ -910,6 +910,11 @@ class MyDict(typing.Dict[T, T]): ...
         class MyDef(typing.DefaultDict[str, T]): ...
         self.assertIs(MyDef[int]().__class__, MyDef)
         self.assertIs(MyDef[int]().__orig_class__, MyDef[int])
+        # ChainMap was added in 3.3
+        if sys.version_info >= (3, 3):
+            class MyChain(typing.ChainMap[str, T]): ...
+            self.assertIs(MyChain[int]().__class__, MyChain)
+            self.assertIs(MyChain[int]().__orig_class__, MyChain[int])
 
     def test_all_repr_eq_any(self):
         objs = (getattr(typing, el) for el in typing.__all__)
@@ -1705,6 +1710,9 @@ def test_deque(self):
         class MyDeque(typing.Deque[int]): ...
         self.assertIsInstance(MyDeque(), collections.deque)
 
+    def test_counter(self):
+        self.assertIsSubclass(collections.Counter, typing.Counter)
+
     def test_set(self):
         self.assertIsSubclass(set, typing.Set)
         self.assertNotIsSubclass(frozenset, typing.Set)
@@ -1772,13 +1780,50 @@ class MyDefDict(typing.DefaultDict[str, int]):
         self.assertIsSubclass(MyDefDict, collections.defaultdict)
         self.assertNotIsSubclass(collections.defaultdict, MyDefDict)
 
+    @skipUnless(sys.version_info >= (3, 3), 'ChainMap was added in 3.3')
+    def test_chainmap_instantiation(self):
+        self.assertIs(type(typing.ChainMap()), collections.ChainMap)
+        self.assertIs(type(typing.ChainMap[KT, VT]()), collections.ChainMap)
+        self.assertIs(type(typing.ChainMap[str, int]()), collections.ChainMap)
+        class CM(typing.ChainMap[KT, VT]): ...
+        self.assertIs(type(CM[int, str]()), CM)
+
+    @skipUnless(sys.version_info >= (3, 3), 'ChainMap was added in 3.3')
+    def test_chainmap_subclass(self):
+
+        class MyChainMap(typing.ChainMap[str, int]):
+            pass
+
+        cm = MyChainMap()
+        self.assertIsInstance(cm, MyChainMap)
+
+        self.assertIsSubclass(MyChainMap, collections.ChainMap)
+        self.assertNotIsSubclass(collections.ChainMap, MyChainMap)
+
     def test_deque_instantiation(self):
         self.assertIs(type(typing.Deque()), collections.deque)
         self.assertIs(type(typing.Deque[T]()), collections.deque)
         self.assertIs(type(typing.Deque[int]()), collections.deque)
         class D(typing.Deque[T]): ...
         self.assertIs(type(D[int]()), D)
 
+    def test_counter_instantiation(self):
+        self.assertIs(type(typing.Counter()), collections.Counter)
+        self.assertIs(type(typing.Counter[T]()), collections.Counter)
+        self.assertIs(type(typing.Counter[int]()), collections.Counter)
+        class C(typing.Counter[T]): ...
+        self.assertIs(type(C[int]()), C)
+
+    def test_counter_subclass_instantiation(self):
+
+        class MyCounter(typing.Counter[int]):
+            pass
+
+        d = MyCounter()
+        self.assertIsInstance(d, MyCounter)
+        self.assertIsInstance(d, typing.Counter)
+        self.assertIsInstance(d, collections.Counter)
+
     def test_no_set_instantiation(self):
         with self.assertRaises(TypeError):
             typing.Set()
diff --git a/src/typing.py b/src/typing.py
@@ -68,6 +68,7 @@
     'SupportsRound',
 
     # Concrete collection types.
+    'Counter',
     'Deque',
     'Dict',
     'DefaultDict',
@@ -1906,6 +1907,31 @@ def __new__(cls, *args, **kwds):
         return _generic_new(collections.defaultdict, cls, *args, **kwds)
 
 
+class Counter(collections.Counter, Dict[T, int], extra=collections.Counter):
+
+    __slots__ = ()
+
+    def __new__(cls, *args, **kwds):
+        if _geqv(cls, Counter):
+            return collections.Counter(*args, **kwds)
+        return _generic_new(collections.Counter, cls, *args, **kwds)
+
+
+if hasattr(collections, 'ChainMap'):
+    # ChainMap only exists in 3.3+
+    __all__.append('ChainMap')
+
+    class ChainMap(collections.ChainMap, MutableMapping[KT, VT],
+                   extra=collections.ChainMap):
+
+        __slots__ = ()
+
+        def __new__(cls, *args, **kwds):
+            if _geqv(cls, ChainMap):
+                return collections.ChainMap(*args, **kwds)
+            return _generic_new(collections.ChainMap, cls, *args, **kwds)
+
+
 # Determine what base class to use for Generator.
 if hasattr(collections_abc, 'Generator'):
     # Sufficiently recent versions of 3.5 have a Generator ABC.
