Fix .keys() and other views to support being iterated more than once

1st1 · 1st1 · commit a9dc2d83794a · 2018-11-20T14:25:07.000-05:00
diff --git a/immutables/_map.c b/immutables/_map.c
@@ -2625,7 +2625,7 @@ map_dump(MapObject *self)
 static int
 map_baseiter_tp_clear(MapIterator *it)
 {
-    Py_CLEAR(it->hi_obj);
+    Py_CLEAR(it->mi_obj);
     return 0;
 }
 
@@ -2640,7 +2640,7 @@ map_baseiter_tp_dealloc(MapIterator *it)
 static int
 map_baseiter_tp_traverse(MapIterator *it, visitproc visit, void *arg)
 {
-    Py_VISIT(it->hi_obj);
+    Py_VISIT(it->mi_obj);
     return 0;
 }
 
@@ -2649,15 +2649,15 @@ map_baseiter_tp_iternext(MapIterator *it)
 {
     PyObject *key;
     PyObject *val;
-    map_iter_t res = map_iterator_next(&it->hi_iter, &key, &val);
+    map_iter_t res = map_iterator_next(&it->mi_iter, &key, &val);
 
     switch (res) {
         case I_END:
             PyErr_SetNone(PyExc_StopIteration);
             return NULL;
 
         case I_ITEM: {
-            return (*(it->hi_yield))(key, val);
+            return (*(it->mi_yield))(key, val);
         }
 
         default: {
@@ -2666,37 +2666,86 @@ map_baseiter_tp_iternext(MapIterator *it)
     }
 }
 
+static int
+map_baseview_tp_clear(MapView *view)
+{
+    Py_CLEAR(view->mv_obj);
+    Py_CLEAR(view->mv_itertype);
+    return 0;
+}
+
+static void
+map_baseview_tp_dealloc(MapView *view)
+{
+    PyObject_GC_UnTrack(view);
+    (void)map_baseview_tp_clear(view);
+    PyObject_GC_Del(view);
+}
+
+static int
+map_baseview_tp_traverse(MapView *view, visitproc visit, void *arg)
+{
+    Py_VISIT(view->mv_obj);
+    return 0;
+}
+
 static Py_ssize_t
-map_baseiter_tp_len(MapIterator *it)
+map_baseview_tp_len(MapView *view)
 {
-    return it->hi_obj->h_count;
+    return view->mv_obj->h_count;
 }
 
-static PyMappingMethods MapIterator_as_mapping = {
-    (lenfunc)map_baseiter_tp_len,
+static PyMappingMethods MapView_as_mapping = {
+    (lenfunc)map_baseview_tp_len,
 };
 
 static PyObject *
-map_baseiter_new(PyTypeObject *type, binaryfunc yield, MapObject *o)
+map_baseview_newiter(PyTypeObject *type, binaryfunc yield, MapObject *map)
 {
-    MapIterator *it = PyObject_GC_New(MapIterator, type);
-    if (it == NULL) {
+    MapIterator *iter = PyObject_GC_New(MapIterator, type);
+    if (iter == NULL) {
+        return NULL;
+    }
+
+    Py_INCREF(map);
+    iter->mi_obj = map;
+    iter->mi_yield = yield;
+    map_iterator_init(&iter->mi_iter, map->h_root);
+
+    PyObject_GC_Track(iter);
+    return (PyObject *)iter;
+}
+
+static PyObject *
+map_baseview_iter(MapView *view)
+{
+    return map_baseview_newiter(
+        view->mv_itertype, view->mv_yield, view->mv_obj);
+}
+
+static PyObject *
+map_baseview_new(PyTypeObject *type, binaryfunc yield,
+                 MapObject *o, PyTypeObject *itertype)
+{
+    MapView *view = PyObject_GC_New(MapView, type);
+    if (view == NULL) {
         return NULL;
     }
 
     Py_INCREF(o);
-    it->hi_obj = o;
-    it->hi_yield = yield;
+    view->mv_obj = o;
+    view->mv_yield = yield;
 
-    map_iterator_init(&it->hi_iter, o->h_root);
+    Py_INCREF(itertype);
+    view->mv_itertype = itertype;
 
-    return (PyObject*)it;
+    PyObject_GC_Track(view);
+    return (PyObject *)view;
 }
 
 #define ITERATOR_TYPE_SHARED_SLOTS                              \
     .tp_basicsize = sizeof(MapIterator),                        \
     .tp_itemsize = 0,                                           \
-    .tp_as_mapping = &MapIterator_as_mapping,                   \
     .tp_dealloc = (destructor)map_baseiter_tp_dealloc,          \
     .tp_getattro = PyObject_GenericGetAttr,                     \
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,        \
@@ -2706,12 +2755,30 @@ map_baseiter_new(PyTypeObject *type, binaryfunc yield, MapObject *o)
     .tp_iternext = (iternextfunc)map_baseiter_tp_iternext,
 
 
+#define VIEW_TYPE_SHARED_SLOTS                                  \
+    .tp_basicsize = sizeof(MapView),                            \
+    .tp_itemsize = 0,                                           \
+    .tp_as_mapping = &MapView_as_mapping,                       \
+    .tp_dealloc = (destructor)map_baseview_tp_dealloc,          \
+    .tp_getattro = PyObject_GenericGetAttr,                     \
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,        \
+    .tp_traverse = (traverseproc)map_baseview_tp_traverse,      \
+    .tp_clear = (inquiry)map_baseview_tp_clear,                 \
+    .tp_iter = (getiterfunc)map_baseview_iter,                  \
+
+
 /////////////////////////////////// _MapItems_Type
 
 
 PyTypeObject _MapItems_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     "items",
+    VIEW_TYPE_SHARED_SLOTS
+};
+
+PyTypeObject _MapItemsIter_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "items_iterator",
     ITERATOR_TYPE_SHARED_SLOTS
 };
 
@@ -2722,10 +2789,11 @@ map_iter_yield_items(PyObject *key, PyObject *val)
 }
 
 static PyObject *
-map_new_iteritems(MapObject *o)
+map_new_items_view(MapObject *o)
 {
-    return map_baseiter_new(
-        &_MapItems_Type, map_iter_yield_items, o);
+    return map_baseview_new(
+        &_MapItems_Type, map_iter_yield_items, o,
+        &_MapItemsIter_Type);
 }
 
 
@@ -2735,6 +2803,12 @@ map_new_iteritems(MapObject *o)
 PyTypeObject _MapKeys_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     "keys",
+    VIEW_TYPE_SHARED_SLOTS
+};
+
+PyTypeObject _MapKeysIter_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "keys_iterator",
     ITERATOR_TYPE_SHARED_SLOTS
 };
 
@@ -2746,19 +2820,32 @@ map_iter_yield_keys(PyObject *key, PyObject *val)
 }
 
 static PyObject *
-map_new_iterkeys(MapObject *o)
+map_new_keys_iter(MapObject *o)
 {
-    return map_baseiter_new(
-        &_MapKeys_Type, map_iter_yield_keys, o);
+    return map_baseview_newiter(
+        &_MapKeysIter_Type, map_iter_yield_keys, o);
 }
 
+static PyObject *
+map_new_keys_view(MapObject *o)
+{
+    return map_baseview_new(
+        &_MapKeys_Type, map_iter_yield_keys, o,
+        &_MapKeysIter_Type);
+}
 
 /////////////////////////////////// _MapValues_Type
 
 
 PyTypeObject _MapValues_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     "values",
+    VIEW_TYPE_SHARED_SLOTS
+};
+
+PyTypeObject _MapValuesIter_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "values_iterator",
     ITERATOR_TYPE_SHARED_SLOTS
 };
 
@@ -2770,10 +2857,11 @@ map_iter_yield_values(PyObject *key, PyObject *val)
 }
 
 static PyObject *
-map_new_itervalues(MapObject *o)
+map_new_values_view(MapObject *o)
 {
-    return map_baseiter_new(
-        &_MapValues_Type, map_iter_yield_values, o);
+    return map_baseview_new(
+        &_MapValues_Type, map_iter_yield_values, o,
+        &_MapValuesIter_Type);
 }
 
 
@@ -2922,7 +3010,7 @@ map_tp_len(BaseMapObject *self)
 static PyObject *
 map_tp_iter(MapObject *self)
 {
-    return map_new_iterkeys(self);
+    return map_new_keys_iter(self);
 }
 
 static PyObject *
@@ -3041,19 +3129,19 @@ map_py_update(MapObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 map_py_items(MapObject *self, PyObject *args)
 {
-    return map_new_iteritems(self);
+    return map_new_items_view(self);
 }
 
 static PyObject *
 map_py_values(MapObject *self, PyObject *args)
 {
-    return map_new_itervalues(self);
+    return map_new_values_view(self);
 }
 
 static PyObject *
 map_py_keys(MapObject *self, PyObject *args)
 {
-    return map_new_iterkeys(self);
+    return map_new_keys_view(self);
 }
 
 static PyObject *
@@ -3844,7 +3932,10 @@ PyInit__map(void)
         (PyType_Ready(&_Map_CollisionNode_Type) < 0) ||
         (PyType_Ready(&_MapKeys_Type) < 0) ||
         (PyType_Ready(&_MapValues_Type) < 0) ||
-        (PyType_Ready(&_MapItems_Type) < 0))
+        (PyType_Ready(&_MapItems_Type) < 0) ||
+        (PyType_Ready(&_MapKeysIter_Type) < 0) ||
+        (PyType_Ready(&_MapValuesIter_Type) < 0) ||
+        (PyType_Ready(&_MapItemsIter_Type) < 0))
     {
         return 0;
     }
diff --git a/immutables/_map.h b/immutables/_map.h
@@ -72,14 +72,25 @@ typedef struct {
    just a key for the 'Keys' iterator, and a value for the 'Values'
    iterator.
 */
+
+typedef struct {
+    PyObject_HEAD
+    MapObject *mv_obj;
+    binaryfunc mv_yield;
+    PyTypeObject *mv_itertype;
+} MapView;
+
 typedef struct {
     PyObject_HEAD
-    MapObject *hi_obj;
-    MapIteratorState hi_iter;
-    binaryfunc hi_yield;
+    MapObject *mi_obj;
+    binaryfunc mi_yield;
+    MapIteratorState mi_iter;
 } MapIterator;
 
 
+/* PyTypes */
+
+
 PyTypeObject _Map_Type;
 PyTypeObject _MapMutation_Type;
 PyTypeObject _Map_ArrayNode_Type;
@@ -88,6 +99,9 @@ PyTypeObject _Map_CollisionNode_Type;
 PyTypeObject _MapKeys_Type;
 PyTypeObject _MapValues_Type;
 PyTypeObject _MapItems_Type;
+PyTypeObject _MapKeysIter_Type;
+PyTypeObject _MapValuesIter_Type;
+PyTypeObject _MapItemsIter_Type;
 
 
 #endif
diff --git a/immutables/map.py b/immutables/map.py
@@ -391,20 +391,43 @@ def dump(self, buf, level):  # pragma: no cover
             buf.append('{}{!r}: {!r}'.format(pad, key, val))
 
 
-class GenWrapper:
+class MapKeys:
 
-    def __init__(self, count, gen):
-        self.__count = count
-        self.__gen = gen
+    def __init__(self, c, m):
+        self.__count = c
+        self.__root = m
+
+    def __len__(self):
+        return self.__count
+
+    def __iter__(self):
+        return iter(self.__root.keys())
+
+
+class MapValues:
+
+    def __init__(self, c, m):
+        self.__count = c
+        self.__root = m
 
     def __len__(self):
         return self.__count
 
     def __iter__(self):
-        return self
+        return iter(self.__root.values())
+
+
+class MapItems:
+
+    def __init__(self, c, m):
+        self.__count = c
+        self.__root = m
 
-    def __next__(self):
-        return next(self.__gen)
+    def __len__(self):
+        return self.__count
+
+    def __iter__(self):
+        return iter(self.__root.items())
 
 
 class Map:
@@ -544,13 +567,13 @@ def __iter__(self):
         yield from self.__root.keys()
 
     def keys(self):
-        return GenWrapper(self.__count, self.__root.keys())
+        return MapKeys(self.__count, self.__root)
 
     def values(self):
-        return GenWrapper(self.__count, self.__root.values())
+        return MapValues(self.__count, self.__root)
 
     def items(self):
-        return GenWrapper(self.__count, self.__root.items())
+        return MapItems(self.__count, self.__root)
 
     def __hash__(self):
         if self.__hash != -1:
diff --git a/tests/test_map.py b/tests/test_map.py
