Patch #403985: Add support for weak-keyed dictionaries

loewis · loewis · commit 5e1633365d8a · 2001-02-27T18:36:56.000Z
diff --git a/Doc/lib/libweakref.tex b/Doc/lib/libweakref.tex
@@ -3,6 +3,8 @@ \section{\module{weakref} ---
 
 \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,23 +20,13 @@ \section{\module{weakref} ---
 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
 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 @@ \section{\module{weakref} ---
   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 @@ \section{\module{weakref} ---
   \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 @@ \subsection{Weak References in Extension Types
 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
@@ -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
@@ -152,6 +152,24 @@ def __repr__(self):
 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
@@ -20,11 +20,14 @@
 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 @@ def values(self):
         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
@@ -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*/
 };
 
 
