blob: 23c22b65ef82a89c0f7171a6b9fbacd0aee3a161 [file] [log] [blame]
Patrick Williamsc124f4f2015-09-15 14:41:29 -05001#
2# This is a copy on write dictionary and set which abuses classes to try and be nice and fast.
3#
Brad Bishopd7bf8c12018-02-25 22:55:05 -05004# Copyright (C) 2006 Tim Ansell
Patrick Williamsc124f4f2015-09-15 14:41:29 -05005#
Andrew Geisslerc9f78652020-09-18 14:11:35 -05006# Please Note:
Patrick Williamsc124f4f2015-09-15 14:41:29 -05007# Be careful when using mutable types (ie Dict and Lists) - operations involving these are SLOW.
8# Assign a file to __warn__ to get warnings about slow operations.
9#
10
Patrick Williamsc0f7c042017-02-23 20:41:17 -060011
Patrick Williamsc124f4f2015-09-15 14:41:29 -050012import copy
Andrew Geisslerc9f78652020-09-18 14:11:35 -050013
Patrick Williamsc124f4f2015-09-15 14:41:29 -050014ImmutableTypes = (
Patrick Williamsc124f4f2015-09-15 14:41:29 -050015 bool,
16 complex,
17 float,
18 int,
Patrick Williamsc124f4f2015-09-15 14:41:29 -050019 tuple,
20 frozenset,
Patrick Williamsc0f7c042017-02-23 20:41:17 -060021 str
Patrick Williamsc124f4f2015-09-15 14:41:29 -050022)
23
24MUTABLE = "__mutable__"
25
Andrew Geisslerc9f78652020-09-18 14:11:35 -050026
Patrick Williamsc124f4f2015-09-15 14:41:29 -050027class COWMeta(type):
28 pass
29
Andrew Geisslerc9f78652020-09-18 14:11:35 -050030
Patrick Williamsc124f4f2015-09-15 14:41:29 -050031class COWDictMeta(COWMeta):
32 __warn__ = False
33 __hasmutable__ = False
34 __marker__ = tuple()
35
36 def __str__(cls):
37 # FIXME: I have magic numbers!
38 return "<COWDict Level: %i Current Keys: %i>" % (cls.__count__, len(cls.__dict__) - 3)
Andrew Geisslerc9f78652020-09-18 14:11:35 -050039
Patrick Williamsc124f4f2015-09-15 14:41:29 -050040 __repr__ = __str__
41
42 def cow(cls):
43 class C(cls):
44 __count__ = cls.__count__ + 1
Andrew Geisslerc9f78652020-09-18 14:11:35 -050045
Patrick Williamsc124f4f2015-09-15 14:41:29 -050046 return C
Andrew Geisslerc9f78652020-09-18 14:11:35 -050047
Patrick Williamsc124f4f2015-09-15 14:41:29 -050048 copy = cow
49 __call__ = cow
50
51 def __setitem__(cls, key, value):
Patrick Williamsc0f7c042017-02-23 20:41:17 -060052 if value is not None and not isinstance(value, ImmutableTypes):
Patrick Williamsc124f4f2015-09-15 14:41:29 -050053 if not isinstance(value, COWMeta):
54 cls.__hasmutable__ = True
55 key += MUTABLE
56 setattr(cls, key, value)
57
58 def __getmutable__(cls, key, readonly=False):
59 nkey = key + MUTABLE
60 try:
61 return cls.__dict__[nkey]
62 except KeyError:
63 pass
64
65 value = getattr(cls, nkey)
66 if readonly:
67 return value
68
69 if not cls.__warn__ is False and not isinstance(value, COWMeta):
70 print("Warning: Doing a copy because %s is a mutable type." % key, file=cls.__warn__)
71 try:
72 value = value.copy()
73 except AttributeError as e:
74 value = copy.copy(value)
75 setattr(cls, nkey, value)
76 return value
77
78 __getmarker__ = []
Andrew Geisslerc9f78652020-09-18 14:11:35 -050079
Patrick Williamsc124f4f2015-09-15 14:41:29 -050080 def __getreadonly__(cls, key, default=__getmarker__):
Andrew Geisslerc9f78652020-09-18 14:11:35 -050081 """
Patrick Williamsc124f4f2015-09-15 14:41:29 -050082 Get a value (even if mutable) which you promise not to change.
83 """
84 return cls.__getitem__(key, default, True)
85
86 def __getitem__(cls, key, default=__getmarker__, readonly=False):
87 try:
88 try:
89 value = getattr(cls, key)
90 except AttributeError:
91 value = cls.__getmutable__(key, readonly)
92
93 # This is for values which have been deleted
94 if value is cls.__marker__:
95 raise AttributeError("key %s does not exist." % key)
96
97 return value
98 except AttributeError as e:
99 if not default is cls.__getmarker__:
100 return default
101
102 raise KeyError(str(e))
103
104 def __delitem__(cls, key):
105 cls.__setitem__(key, cls.__marker__)
106
107 def __revertitem__(cls, key):
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600108 if key not in cls.__dict__:
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500109 key += MUTABLE
110 delattr(cls, key)
111
112 def __contains__(cls, key):
113 return cls.has_key(key)
114
115 def has_key(cls, key):
116 value = cls.__getreadonly__(key, cls.__marker__)
117 if value is cls.__marker__:
118 return False
119 return True
120
121 def iter(cls, type, readonly=False):
122 for key in dir(cls):
123 if key.startswith("__"):
124 continue
125
126 if key.endswith(MUTABLE):
127 key = key[:-len(MUTABLE)]
128
129 if type == "keys":
130 yield key
131
132 try:
133 if readonly:
134 value = cls.__getreadonly__(key)
135 else:
136 value = cls[key]
137 except KeyError:
138 continue
139
140 if type == "values":
141 yield value
142 if type == "items":
143 yield (key, value)
Brad Bishop1a4b7ee2018-12-16 17:11:34 -0800144 return
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500145
146 def iterkeys(cls):
147 return cls.iter("keys")
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500148
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500149 def itervalues(cls, readonly=False):
150 if not cls.__warn__ is False and cls.__hasmutable__ and readonly is False:
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500151 print("Warning: If you aren't going to change any of the values call with True.", file=cls.__warn__)
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500152 return cls.iter("values", readonly)
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500153
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500154 def iteritems(cls, readonly=False):
155 if not cls.__warn__ is False and cls.__hasmutable__ and readonly is False:
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500156 print("Warning: If you aren't going to change any of the values call with True.", file=cls.__warn__)
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500157 return cls.iter("items", readonly)
158
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500159
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500160class COWSetMeta(COWDictMeta):
161 def __str__(cls):
162 # FIXME: I have magic numbers!
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500163 return "<COWSet Level: %i Current Keys: %i>" % (cls.__count__, len(cls.__dict__) - 3)
164
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500165 __repr__ = __str__
166
167 def cow(cls):
168 class C(cls):
169 __count__ = cls.__count__ + 1
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500170
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500171 return C
172
173 def add(cls, value):
174 COWDictMeta.__setitem__(cls, repr(hash(value)), value)
175
176 def remove(cls, value):
177 COWDictMeta.__delitem__(cls, repr(hash(value)))
178
179 def __in__(cls, value):
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600180 return repr(hash(value)) in COWDictMeta
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500181
182 def iterkeys(cls):
183 raise TypeError("sets don't have keys")
184
185 def iteritems(cls):
186 raise TypeError("sets don't have 'items'")
187
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500188
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500189# These are the actual classes you use!
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500190class COWDictBase(metaclass=COWDictMeta):
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500191 __count__ = 0
192
Andrew Geisslerc9f78652020-09-18 14:11:35 -0500193
194class COWSetBase(metaclass=COWSetMeta):
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500195 __count__ = 0