Patch #403985: Add support for weak-keyed dictionaries

5 files changed, 158 insertions, 21 deletions

diff --git a/Doc/lib/libweakref.tex b/Doc/lib/libweakref.tex
index 61ed9fe..97f624d 100644
--- a/Doc/lib/libweakref.tex
+++ b/Doc/lib/libweakref.tex
@@ -3,6 +3,8 @@
 
 \declaremodule{extension}{weakref}
 \moduleauthor{Fred L. Drake, Jr.}{fdrake@acm.org}
+\moduleauthor{Neil Schemenauer}{nas@arctrix.com}
+\moduleauthor{Martin von L\o"wis}{martin@loewis.home.cs.tu-berlin.de}
 \sectionauthor{Fred L. Drake, Jr.}{fdrake@acm.org}
 
 \versionadded{2.1}
@@ -18,13 +20,6 @@ include class instances and dictionaries.  Extension types can easily
 be made to support weak references; see section \ref{weakref-extension},
 ``Weak References in Extension Types,'' for more information.
 
-
-\strong{Warning:}
-The weak dictionaries provided in the current implementation and
-described below are subject to change.  They are included to solicit
-feedback and usage experience, and may be changed or removed in the
-final version.
-
 \strong{Warning:}
 The desired semantics of weak-reference proxy objects are not
 completely clear; it is very difficult to create proxies which behave
@@ -32,9 +27,6 @@ exactly like the type of the referent.  The details of these objects
 are likely to change to some degree before the final release as
 experience reports become available.
 
-Please send specific feedback on this module to Fred Drake at
-\email{fdrake@acm.org}.
-
 
 \begin{funcdesc}{ref}{object\optional{, callback}}
   Return a weak reference to \var{object}.  If \var{callback} is
@@ -53,15 +45,36 @@ Please send specific feedback on this module to Fred Drake at
   error output, but cannot be propagated; they are handled in exactly
   the same way as exceptions raised from an object's
   \method{__del__()} method.
+  
+  Weak references are hashable if the \var{object} is hashable.  They
+  will maintain their hash value even after the \var{object} was
+  deleted.  If \function{hash()} is called the first time only after
+  the \var{object} was deleted, the call will raise
+  \exception{TypeError}.
+  
+  Weak references support test for equality, but not ordering.  If the
+  \var{object} is still alive, to references are equal if the objects
+  are equal (regardless of the \var{callback}).  If the \var{object}
+  has been deleted, they are equal iff they are identical.
+
 \end{funcdesc}
 
-\begin{funcdesc}{mapping}{\optional{dict}}
+\begin{funcdesc}{mapping}{\optional{dict\optional{, weakkeys=0}}}
   Return a weak dictionary.  If \var{dict} is given and not
   \code{None}, the new dictionary will contain the items contained in
   \var{dict}.  The values from \var{dict} must be weakly referencable;
   if any values which would be inserted into the new mapping are not
   weakly referencable, \exception{TypeError} will be raised and the
   new mapping will be empty.
+
+  If the \var{weakkeys} argument is not given or zero, the values in
+  the dictionary are weak. That means the entries in the dictionary
+  will be discarded when no strong reference to the value exists
+  anymore.
+
+  If the \var{weakkeys} argument is nonzero, the keys in the
+  dictionary are weak, i.e. the entry in the dictionary is discarded
+  when the last strong reference to the key is discarded.
 \end{funcdesc}
 
 \begin{funcdesc}{proxy}{object\optional{, callback}}
@@ -87,9 +100,16 @@ Please send specific feedback on this module to Fred Drake at
   \var{object}.
 \end{funcdesc}
 
-\begin{classdesc}{WeakDictionary}{\optional{dict}}
-  The class of the mapping objects returned by \function{mapping()}.
-  This can be used for subclassing the implementation if needed.
+\begin{classdesc}{WeakKeyDictionary}{\optional{dict}}
+  The class of the mapping objects returned by \function{mapping()}
+  when \var{weakkeys} is true.  This can be used for subclassing the
+  implementation if needed.
+\end{classdesc}
+
+\begin{classdesc}{WeakValueDictionary}{\optional{dict}}
+  The class of the mapping objects returned by \function{mapping()}
+  when \var{weakkeys} if false.  This can be used for subclassing the
+  implementation if needed.
 \end{classdesc}
 
 \begin{datadesc}{ReferenceType}
@@ -187,8 +207,8 @@ For example, the instance type is defined with the following structure:
 typedef struct {
     PyObject_HEAD
     PyClassObject *in_class;       /* The class object */
-    PyObject	  *in_dict;        /* A dictionary */
-    PyObject	  *in_weakreflist; /* List of weak references */
+    PyObject      *in_dict;        /* A dictionary */
+    PyObject      *in_weakreflist; /* List of weak references */
 } PyInstanceObject;
 \end{verbatim}
 
diff --git a/Lib/test/output/test_weakref b/Lib/test/output/test_weakref
index b3d6f97..ef74e7e 100644
--- a/Lib/test/output/test_weakref
+++ b/Lib/test/output/test_weakref
@@ -15,6 +15,10 @@ Weak Valued Dictionaries
 objects are stored in weak dict
 weak dict test complete
 
+Weak Keyed Dictionaries
+objects are stored in weak dict
+weak key dict test complete
+
 Non-callable Proxy References
 XXX -- tests not written!
 
diff --git a/Lib/test/test_weakref.py b/Lib/test/test_weakref.py
index fe19373..befa70d 100644
--- a/Lib/test/test_weakref.py
+++ b/Lib/test/test_weakref.py
@@ -153,6 +153,24 @@ dict.clear()
 print "weak dict test complete"
 
 print
+print "Weak Keyed Dictionaries"
+
+dict = weakref.mapping(weakkeys=1)
+objects = map(Object, range(10))
+for o in objects:
+    dict[o] = o.arg
+print "objects are stored in weak dict"
+for o in objects:
+    verify(weakref.getweakrefcount(o) == 1,
+           "wrong number of weak references to %r!" % o)
+    verify(o.arg is dict[o],
+           "wrong object returned by weak dict!")
+del objects,o
+verify(len(dict)==0, "deleting the keys did not clear the dictionary")
+print "weak key dict test complete"
+
+
+print
 print "Non-callable Proxy References"
 print "XXX -- tests not written!"
 
diff --git a/Lib/weakref.py b/Lib/weakref.py
index f1222fa..9d5eac0 100644
--- a/Lib/weakref.py
+++ b/Lib/weakref.py
@@ -20,11 +20,14 @@ from _weakref import \
 ProxyTypes = (ProxyType, CallableProxyType)
 
 
-def mapping(dict=None):
-    return WeakDictionary(dict)
+def mapping(dict=None,weakkeys=0):
+    if weakkeys:
+        return WeakKeyDictionary(dict)
+    else:
+        return WeakValueDictionary(dict)
 
 
-class WeakDictionary(UserDict.UserDict):
+class WeakValueDictionary(UserDict.UserDict):
 
     # We inherit the constructor without worrying about the input
     # dictionary; since it uses our .update() method, we get the right
@@ -112,5 +115,59 @@ class WeakDictionary(UserDict.UserDict):
         return L
 
 
+class WeakKeyDictionary(UserDict.UserDict):
+
+    def __init__(self, dict=None):
+        self.data = {}
+        if dict is not None: self.update(dict)
+        def remove(k, data=self.data):
+            del data[k]
+        self._remove = remove
+
+    def __getitem__(self, key):
+        return self.data[ref(key)]
+
+    def __repr__(self):
+        return "<WeakKeyDictionary at %s>" % id(self)
+
+    def __setitem__(self, key, value):
+        self.data[ref(key, self._remove)] = value
+
+    def copy(self):
+        new = WeakKeyDictionary()
+        for key, value in self.data.items():
+            o = key()
+            if o is not None:
+                new[o] = value
+
+    def get(self, key, default):
+        return self.data.get(ref(key),default)
+
+    def items(self):
+        L = []
+        for key, value in self.data.items():
+            o = key()
+            if o is not None:
+                L.append((o, value))
+        return L
+
+    def popitem(self):
+        while 1:
+            key, value = self.data.popitem()
+            o = key()
+            if o is not None:
+                return o, value
+
+    def setdefault(self, key, default):
+        return self.data.setdefault(ref(key, self._remove),default)
+
+    def update(self, dict):
+        d = self.data
+        L = []
+        for key, value in dict.items():
+            L.append(ref(key, self._remove), value)
+        for key, r in L:
+            d[key] = r
+
 # no longer needed
 del UserDict
diff --git a/Modules/_weakref.c b/Modules/_weakref.c
index fbf886a..399b4fe 100644
--- a/Modules/_weakref.c
+++ b/Modules/_weakref.c
@@ -8,6 +8,7 @@ struct _PyWeakReference {
     PyObject_HEAD
     PyObject *wr_object;
     PyObject *wr_callback;
+    long hash;
     PyWeakReference *wr_prev;
     PyWeakReference *wr_next;
 };
@@ -39,6 +40,8 @@ new_weakref(void)
     else {
         result = PyObject_NEW(PyWeakReference, &PyWeakReference_Type);
     }
+    if (result)
+        result->hash = -1;
     return result;
 }
 
@@ -112,6 +115,20 @@ weakref_call(PyWeakReference *self, PyObject *args, PyObject *kw)
 }
 
 
+static long
+weakref_hash(PyWeakReference *self)
+{
+    if (self->hash != -1)
+        return self->hash;
+    if (self->wr_object == Py_None) {
+        PyErr_SetString(PyExc_TypeError, "weak object has gone away");
+        return -1;
+    }
+    self->hash = PyObject_Hash(self->wr_object);
+    return self->hash;
+}
+    
+
 static PyObject *
 weakref_repr(PyWeakReference *self)
 {
@@ -128,6 +145,25 @@ weakref_repr(PyWeakReference *self)
     return PyString_FromString(buffer);
 }
 
+/* Weak references only support equality, not ordering. Two weak references
+   are equal if the underlying objects are equal. If the underlying object has
+   gone away, they are equal if they are identical. */
+
+static PyObject *
+weakref_richcompare(PyWeakReference* self, PyWeakReference* other, int op)
+{
+    if (op != Py_EQ || self->ob_type != other->ob_type) {
+        Py_INCREF(Py_NotImplemented);
+        return Py_NotImplemented;
+    }
+    if (self->wr_object == Py_None || other->wr_object == Py_None) {
+        PyObject *res = self==other ? Py_True : Py_False;
+        Py_INCREF(res);
+        return res;
+    }
+    return PyObject_RichCompare(self->wr_object, other->wr_object, op);
+}
+
 
 statichere PyTypeObject
 PyWeakReference_Type = {
@@ -145,16 +181,18 @@ PyWeakReference_Type = {
     0,                          /*tp_as_number*/
     0,                          /*tp_as_sequence*/
     0,                          /*tp_as_mapping*/
-    0,                          /*tp_hash*/
+    (hashfunc)weakref_hash,      /*tp_hash*/
     (ternaryfunc)weakref_call,  /*tp_call*/
     0,                          /*tp_str*/
     0,                          /*tp_getattro*/
     0,                          /*tp_setattro*/
     0,                          /*tp_as_buffer*/
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_GC,
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_GC | Py_TPFLAGS_HAVE_RICHCOMPARE,
     0,                          /*tp_doc*/
     (traverseproc)gc_traverse,  /*tp_traverse*/
     (inquiry)gc_clear,          /*tp_clear*/
+    (richcmpfunc)weakref_richcompare,	/*tp_richcompare*/
+    0,				/*tp_weaklistoffset*/
 };