Adding Generics to Pyrsistent Classes

pyrsistent/_checked_types.py

@@ -2,11 +2,16 @@
 
 from abc import abstractmethod, ABCMeta
 from collections.abc import Iterable
+from typing import TypeVar, Generic
 
 from pyrsistent._pmap import PMap, pmap
 from pyrsistent._pset import PSet, pset
 from pyrsistent._pvector import PythonPVector, python_pvector
 
+T_co = TypeVar('T_co', covariant=True)
+KT = TypeVar('KT')
+VT_co = TypeVar('VT_co', covariant=True)
+
 
 class CheckedType(object):
     """
@@ -271,7 +276,7 @@ def _checked_type_create(cls, source_data, _factory_fields=None, ignore_extra=Fa
 
     return cls(source_data)
 
-class CheckedPVector(PythonPVector, CheckedType, metaclass=_CheckedTypeMeta):
+class CheckedPVector(Generic[T_co], PythonPVector, CheckedType, metaclass=_CheckedTypeMeta):
     """
     A CheckedPVector is a PVector which allows specifying type and invariant checks.
 
@@ -357,7 +362,7 @@ def evolver(self):
         return CheckedPVector.Evolver(self.__class__, self)
 
 
-class CheckedPSet(PSet, CheckedType, metaclass=_CheckedTypeMeta):
+class CheckedPSet(PSet[T_co], CheckedType, metaclass=_CheckedTypeMeta):
     """
     A CheckedPSet is a PSet which allows specifying type and invariant checks.
 
@@ -455,7 +460,7 @@ def default_serializer(self, _, key, value):
 _UNDEFINED_CHECKED_PMAP_SIZE = object()
 
 
-class CheckedPMap(PMap, CheckedType, metaclass=_CheckedMapTypeMeta):
+class CheckedPMap(PMap[KT, VT_co], CheckedType, metaclass=_CheckedMapTypeMeta):
     """
     A CheckedPMap is a PMap which allows specifying type and invariant checks.
 

pyrsistent/_pbag.py

@@ -1,13 +1,16 @@
 from collections.abc import Container, Iterable, Sized, Hashable
 from functools import reduce
+from typing import Generic, TypeVar
 from pyrsistent._pmap import pmap
 
+T_co = TypeVar('T_co', covariant=True)
+
 
 def _add_to_counters(counters, element):
     return counters.set(element, counters.get(element, 0) + 1)
 
 
-class PBag(object):
+class PBag(Generic[T_co]):
     """
     A persistent bag/multiset type.
 

pyrsistent/_pdeque.py

@@ -1,10 +1,13 @@
 from collections.abc import Sequence, Hashable
 from itertools import islice, chain
 from numbers import Integral
+from typing import TypeVar, Generic
 from pyrsistent._plist import plist
 
+T = TypeVar("T")
 
-class PDeque(object):
+
+class PDeque(Generic[T]):
     """
     Persistent double ended queue (deque). Allows quick appends and pops in both ends. Implemented
     using two persistent lists.
@@ -175,7 +178,7 @@ def __eq__(self, other):
         return False
 
     def __hash__(self):
-        return  hash(tuple(self))
+        return hash(tuple(self))
 
     def __len__(self):
         return self._length
@@ -275,7 +278,7 @@ def remove(self, elem):
             try:
                 # This is severely inefficient with a double reverse, should perhaps implement a remove_last()?
                 return PDeque(self._left_list,
-                               self._right_list.reverse().remove(elem).reverse(), self._length - 1)
+                              self._right_list.reverse().remove(elem).reverse(), self._length - 1)
             except ValueError as e:
                 raise ValueError('{0} not found in PDeque'.format(elem)) from e
 

pyrsistent/_plist.py

@@ -1,6 +1,9 @@
 from collections.abc import Sequence, Hashable
 from numbers import Integral
 from functools import reduce
+from typing import Generic, TypeVar
+
+T_co = TypeVar('T_co', covariant=True)
 
 
 class _PListBuilder(object):
@@ -219,7 +222,7 @@ def remove(self, elem):
         raise ValueError('{0} not found in PList'.format(elem))
 
 
-class PList(_PListBase):
+class PList(Generic[T_co], _PListBase):
     """
     Classical Lisp style singly linked list. Adding elements to the head using cons is O(1).
     Element access is O(k) where k is the position of the element in the list. Taking the

pyrsistent/_pmap.py

@@ -1,8 +1,12 @@
 from collections.abc import Mapping, Hashable
 from itertools import chain
+from typing import Generic, TypeVar
+
 from pyrsistent._pvector import pvector
 from pyrsistent._transformations import transform
 
+KT = TypeVar('KT')
+VT_co = TypeVar('VT_co', covariant=True)
 class PMapView:
     """View type for the persistent map/dict type `PMap`.
 
@@ -103,7 +107,7 @@ def __eq__(self, x):
         elif not isinstance(x, type(self)): return False
         else: return self._map == x._map
 
-class PMap(object):
+class PMap(Generic[KT, VT_co]):
     """
     Persistent map/dict. Tries to follow the same naming conventions as the built in dict where feasible.
 

pyrsistent/_pset.py

@@ -1,9 +1,12 @@
 from collections.abc import Set, Hashable
 import sys
+from typing import TypeVar, Generic
 from pyrsistent._pmap import pmap
 
+T_co = TypeVar('T_co', covariant=True)
 
-class PSet(object):
+
+class PSet(Generic[T_co]):
     """
     Persistent set implementation. Built on top of the persistent map. The set supports all operations
     in the Set protocol and is Hashable.

pyrsistent/_pvector.py

@@ -2,8 +2,12 @@
 from collections.abc import Sequence, Hashable
 from numbers import Integral
 import operator
+from typing import TypeVar, Generic
+
 from pyrsistent._transformations import transform
 
+T_co = TypeVar('T_co', covariant=True)
+
 
 def _bitcount(val):
     return bin(val).count("1")
@@ -410,7 +414,7 @@ def remove(self, value):
         l.remove(value)
         return _EMPTY_PVECTOR.extend(l)
 
-class PVector(metaclass=ABCMeta):
+class PVector(Generic[T_co],metaclass=ABCMeta):
     """
     Persistent vector implementation. Meant as a replacement for the cases where you would normally
     use a Python list.