Ticket #13394: weak_value_dictionary.pyx

File weak_value_dictionary.pyx, 5.2 KB (added by SimonKing, 5 years ago)

Proof of concept

Line 
1import weakref
2from weakref import KeyedRef
3import sage
4from sage.all import add
5
6from cpython.dict cimport *
7from cpython.weakref cimport *
8
9cdef PyObject* Py_None = <PyObject*>None
10
11cdef class WVD(dict):
12    def __init__(self, data=()):
13        dict.__init__(self)
14        for k,v in data:
15            self[k] = v
16    def __getitem__(self, k):
17        cdef void * buckref = PyDict_GetItem(self, hash(k))
18        if buckref==NULL:
19            raise KeyError(k)
20        cdef list bucket = <list>buckref
21        cdef PyObject* out
22        for K,v in bucket:
23            if K==k:
24                out = PyWeakref_GetObject(v)
25                if out!=Py_None:
26                    return <object>out
27                break
28        raise KeyError(k)
29
30    def __repr__(self):
31        return "<WeakValueDictionary at %s>" % id(self)
32
33    def get(self, k, default=None):
34        cdef void * buckref = PyDict_GetItem(self, hash(k))
35        if buckref==NULL:
36            return default
37        cdef list bucket = <list>buckref
38        cdef PyObject* out
39        for K,v in bucket:
40            if K==k:
41                out = PyWeakref_GetObject(v)
42                if out!=Py_None:
43                    return <object>out
44                break
45        return default
46
47    def callback(self, r):
48        hashk,idk = r.key
49        cdef void * buckref = PyDict_GetItem(self, hashk)
50        if buckref==NULL:
51            return
52        cdef list bucket = <list>buckref
53        cdef int i,l
54        l = len(bucket)
55        for i from 0 <= i < l:
56            if id((<tuple>(bucket[i]))[0]) == idk:
57                del bucket[i]
58                if not bucket:
59                    PyDict_DelItem(self,hashk)
60                return
61
62    def setdefault(self, k, default=None):
63        cdef long hashk = hash(k)
64        cdef void * buckref = PyDict_GetItem(self, hashk)
65        cdef list bucket
66        cdef object k0
67        cdef int i,l
68        cdef PyObject* out
69        if buckref==NULL:
70            bucket = [(k, KeyedRef(default, self.callback, (hash(k),id(k))))]
71            PyDict_SetItem(self,hashk,bucket)
72            return default
73        bucket = <list>buckref
74        l = len(bucket)
75        for i from 0 <= i < l:
76            k0, v = <tuple>(bucket[i])
77            if k0 == k:
78                out = PyWeakref_GetObject(v)
79                if out!=Py_None:
80                    return <object>out
81                bucket[i] = (k, KeyedRef(default, self.callback, (hashk,id(k))))
82                return default
83        bucket.append((k, KeyedRef(default, self.callback, (hashk,id(k)))))
84        return default
85
86    def __setitem__(self, k, v):
87        cdef long hashk = hash(k)
88        cdef void * buckref = PyDict_GetItem(self, hashk)
89        cdef list bucket
90        if buckref==NULL:
91            bucket = []
92            PyDict_SetItem(self,hashk,bucket)
93        else:
94            bucket = <list>buckref
95        cdef object k0
96        cdef int i,l
97        l = len(bucket)
98        for i from 0 <= i < l:
99            k0 = <tuple>(bucket[i])[0]
100            if k0 == k:
101                bucket[i] = (k, KeyedRef(v, self.callback, (hashk,id(k))))
102                return
103        bucket.append((k, KeyedRef(v, self.callback, (hashk,id(k)))))
104
105    def __delitem__(self, k):
106        cdef long hashk = hash(k)
107        cdef void * buckref = PyDict_GetItem(self, hashk)
108        if buckref==NULL:
109            raise KeyError(k)
110        cdef list bucket = <list>buckref
111        cdef int i,l
112        l = len(bucket)
113        for i from 0 <= i < l:
114            if <tuple>(bucket[i])[0] == k:
115                del bucket[i]
116                if not bucket:
117                    PyDict_DelItem(self,hashk)
118                return
119
120    def __contains__(self, k):
121        cdef list bucket
122        cdef void * buckref = PyDict_GetItem(self, hash(k))
123        if buckref==NULL:
124            return False
125        bucket = <list>buckref
126        cdef int i,l
127        l = len(bucket)
128        for i from 0 <= i < l:
129            K,v = <tuple>(bucket[i])
130            if K == k:
131                if PyWeakref_GetObject(v)!=Py_None:
132                    return True
133                else:
134                    return False
135        return False
136
137    def __len__(self):
138        return add(len(<object>PyDict_GetItem(self,basekey)) for basekey in (<dict>self).iterkeys())
139
140    def iterkeys(self):
141        cdef list bucket
142        for basekey in (<dict>self).iterkeys():
143            bucket = <list>PyDict_GetItem(self, basekey)
144            for k,v in bucket:
145                if PyWeakref_GetObject(v)!=Py_None:
146                    yield k
147
148    def __iter__(self):
149        return self.iterkeys()
150
151    def itervalues(self):
152        cdef PyObject * obj
153        cdef list bucket
154        for basekey in (<dict>self).iterkeys():
155            bucket = <list>PyDict_GetItem(self, basekey)
156            for k,v in bucket:
157                obj = PyWeakref_GetObject(v)
158                if obj != Py_None:
159                    yield <object>obj
160
161    def iteritems(self):
162        cdef PyObject * obj
163        cdef list bucket
164        for basekey in (<dict>self).iterkeys():
165            bucket = <list>PyDict_GetItem(self, basekey)
166            for k,v in bucket:
167                obj = PyWeakref_GetObject(v)
168                if obj != Py_None:
169                    yield k, <object>obj