| 1 | # (c) 2005 Ian Bicking and contributors; written for Paste (http://pythonpaste.org) |
|---|
| 2 | # Licensed under the MIT license: http://www.opensource.org/licenses/mit-license.php |
|---|
| 3 | import cgi |
|---|
| 4 | import copy |
|---|
| 5 | import sys |
|---|
| 6 | from UserDict import DictMixin |
|---|
| 7 | |
|---|
| 8 | class MultiDict(DictMixin): |
|---|
| 9 | |
|---|
| 10 | """ |
|---|
| 11 | An ordered dictionary that can have multiple values for each key. |
|---|
| 12 | Adds the methods getall, getone, mixed, and add to the normal |
|---|
| 13 | dictionary interface. |
|---|
| 14 | """ |
|---|
| 15 | |
|---|
| 16 | def __init__(self, *args, **kw): |
|---|
| 17 | if len(args) > 1: |
|---|
| 18 | raise TypeError( |
|---|
| 19 | "MultiDict can only be called with one positional argument") |
|---|
| 20 | if args: |
|---|
| 21 | if hasattr(args[0], 'iteritems'): |
|---|
| 22 | items = list(args[0].iteritems()) |
|---|
| 23 | elif hasattr(args[0], 'items'): |
|---|
| 24 | items = args[0].items() |
|---|
| 25 | else: |
|---|
| 26 | items = list(args[0]) |
|---|
| 27 | self._items = items |
|---|
| 28 | else: |
|---|
| 29 | self._items = [] |
|---|
| 30 | self._items.extend(kw.iteritems()) |
|---|
| 31 | |
|---|
| 32 | def __getitem__(self, key): |
|---|
| 33 | for k, v in self._items: |
|---|
| 34 | if k == key: |
|---|
| 35 | return v |
|---|
| 36 | raise KeyError(repr(key)) |
|---|
| 37 | |
|---|
| 38 | def __setitem__(self, key, value): |
|---|
| 39 | try: |
|---|
| 40 | del self[key] |
|---|
| 41 | except KeyError: |
|---|
| 42 | pass |
|---|
| 43 | self._items.append((key, value)) |
|---|
| 44 | |
|---|
| 45 | def add(self, key, value): |
|---|
| 46 | """ |
|---|
| 47 | Add the key and value, not overwriting any previous value. |
|---|
| 48 | """ |
|---|
| 49 | self._items.append((key, value)) |
|---|
| 50 | |
|---|
| 51 | def getall(self, key): |
|---|
| 52 | """ |
|---|
| 53 | Return a list of all values matching the key (may be an empty list) |
|---|
| 54 | """ |
|---|
| 55 | result = [] |
|---|
| 56 | for k, v in self._items: |
|---|
| 57 | if key == k: |
|---|
| 58 | result.append(v) |
|---|
| 59 | return result |
|---|
| 60 | |
|---|
| 61 | def getone(self, key): |
|---|
| 62 | """ |
|---|
| 63 | Get one value matching the key, raising a KeyError if multiple |
|---|
| 64 | values were found. |
|---|
| 65 | """ |
|---|
| 66 | v = self.getall(key) |
|---|
| 67 | if not v: |
|---|
| 68 | raise KeyError('Key not found: %r' % key) |
|---|
| 69 | if len(v) > 1: |
|---|
| 70 | raise KeyError('Multiple values match %r: %r' % (key, v)) |
|---|
| 71 | return v[0] |
|---|
| 72 | |
|---|
| 73 | def mixed(self): |
|---|
| 74 | """ |
|---|
| 75 | Returns a dictionary where the values are either single |
|---|
| 76 | values, or a list of values when a key/value appears more than |
|---|
| 77 | once in this dictionary. This is similar to the kind of |
|---|
| 78 | dictionary often used to represent the variables in a web |
|---|
| 79 | request. |
|---|
| 80 | """ |
|---|
| 81 | result = {} |
|---|
| 82 | multi = {} |
|---|
| 83 | for key, value in self._items: |
|---|
| 84 | if key in result: |
|---|
| 85 | # We do this to not clobber any lists that are |
|---|
| 86 | # *actual* values in this dictionary: |
|---|
| 87 | if key in multi: |
|---|
| 88 | result[key].append(value) |
|---|
| 89 | else: |
|---|
| 90 | result[key] = [result[key], value] |
|---|
| 91 | multi[key] = None |
|---|
| 92 | else: |
|---|
| 93 | result[key] = value |
|---|
| 94 | return result |
|---|
| 95 | |
|---|
| 96 | def dict_of_lists(self): |
|---|
| 97 | """ |
|---|
| 98 | Returns a dictionary where each key is associated with a |
|---|
| 99 | list of values. |
|---|
| 100 | """ |
|---|
| 101 | result = {} |
|---|
| 102 | for key, value in self._items: |
|---|
| 103 | if key in result: |
|---|
| 104 | result[key].append(value) |
|---|
| 105 | else: |
|---|
| 106 | result[key] = [value] |
|---|
| 107 | return result |
|---|
| 108 | |
|---|
| 109 | def __delitem__(self, key): |
|---|
| 110 | items = self._items |
|---|
| 111 | found = False |
|---|
| 112 | for i in range(len(items)-1, -1, -1): |
|---|
| 113 | if items[i][0] == key: |
|---|
| 114 | del items[i] |
|---|
| 115 | found = True |
|---|
| 116 | if not found: |
|---|
| 117 | raise KeyError(repr(key)) |
|---|
| 118 | |
|---|
| 119 | def __contains__(self, key): |
|---|
| 120 | for k, v in self._items: |
|---|
| 121 | if k == key: |
|---|
| 122 | return True |
|---|
| 123 | return False |
|---|
| 124 | |
|---|
| 125 | has_key = __contains__ |
|---|
| 126 | |
|---|
| 127 | def clear(self): |
|---|
| 128 | self._items = [] |
|---|
| 129 | |
|---|
| 130 | def copy(self): |
|---|
| 131 | return MultiDict(self) |
|---|
| 132 | |
|---|
| 133 | def setdefault(self, key, default=None): |
|---|
| 134 | for k, v in self._items: |
|---|
| 135 | if key == k: |
|---|
| 136 | return v |
|---|
| 137 | self._items.append((key, default)) |
|---|
| 138 | return default |
|---|
| 139 | |
|---|
| 140 | def pop(self, key, *args): |
|---|
| 141 | if len(args) > 1: |
|---|
| 142 | raise TypeError, "pop expected at most 2 arguments, got "\ |
|---|
| 143 | + repr(1 + len(args)) |
|---|
| 144 | for i in range(len(self._items)): |
|---|
| 145 | if self._items[i][0] == key: |
|---|
| 146 | v = self._items[i][1] |
|---|
| 147 | del self._items[i] |
|---|
| 148 | return v |
|---|
| 149 | if args: |
|---|
| 150 | return args[0] |
|---|
| 151 | else: |
|---|
| 152 | raise KeyError(repr(key)) |
|---|
| 153 | |
|---|
| 154 | def popitem(self): |
|---|
| 155 | return self._items.pop() |
|---|
| 156 | |
|---|
| 157 | def update(self, other=None, **kwargs): |
|---|
| 158 | if other is None: |
|---|
| 159 | pass |
|---|
| 160 | elif hasattr(other, 'items'): |
|---|
| 161 | self._items.extend(other.items()) |
|---|
| 162 | elif hasattr(other, 'keys'): |
|---|
| 163 | for k in other.keys(): |
|---|
| 164 | self._items.append((k, other[k])) |
|---|
| 165 | else: |
|---|
| 166 | for k, v in other: |
|---|
| 167 | self._items.append((k, v)) |
|---|
| 168 | if kwargs: |
|---|
| 169 | self.update(kwargs) |
|---|
| 170 | |
|---|
| 171 | def __repr__(self): |
|---|
| 172 | items = ', '.join(['(%r, %r)' % v for v in self._items]) |
|---|
| 173 | return '%s([%s])' % (self.__class__.__name__, items) |
|---|
| 174 | |
|---|
| 175 | def __len__(self): |
|---|
| 176 | return len(self._items) |
|---|
| 177 | |
|---|
| 178 | ## |
|---|
| 179 | ## All the iteration: |
|---|
| 180 | ## |
|---|
| 181 | |
|---|
| 182 | def keys(self): |
|---|
| 183 | return [k for k, v in self._items] |
|---|
| 184 | |
|---|
| 185 | def iterkeys(self): |
|---|
| 186 | for k, v in self._items: |
|---|
| 187 | yield k |
|---|
| 188 | |
|---|
| 189 | __iter__ = iterkeys |
|---|
| 190 | |
|---|
| 191 | def items(self): |
|---|
| 192 | return self._items[:] |
|---|
| 193 | |
|---|
| 194 | def iteritems(self): |
|---|
| 195 | return iter(self._items) |
|---|
| 196 | |
|---|
| 197 | def values(self): |
|---|
| 198 | return [v for k, v in self._items] |
|---|
| 199 | |
|---|
| 200 | def itervalues(self): |
|---|
| 201 | for k, v in self._items: |
|---|
| 202 | yield v |
|---|
| 203 | |
|---|
| 204 | class UnicodeMultiDict(DictMixin): |
|---|
| 205 | """ |
|---|
| 206 | A MultiDict wrapper that decodes returned values to unicode on the |
|---|
| 207 | fly. Decoding is not applied to assigned values. |
|---|
| 208 | |
|---|
| 209 | The key/value contents are assumed to be ``str``/``strs`` or |
|---|
| 210 | ``str``/``FieldStorages`` (as is returned by the ``paste.request.parse_`` |
|---|
| 211 | functions). |
|---|
| 212 | |
|---|
| 213 | Can optionally also decode keys when the ``decode_keys`` argument is |
|---|
| 214 | True. |
|---|
| 215 | |
|---|
| 216 | ``FieldStorage`` instances are cloned, and the clone's ``filename`` |
|---|
| 217 | variable is decoded. Its ``name`` variable is decoded when ``decode_keys`` |
|---|
| 218 | is enabled. |
|---|
| 219 | |
|---|
| 220 | """ |
|---|
| 221 | def __init__(self, multi=None, encoding=None, errors='strict', |
|---|
| 222 | decode_keys=False): |
|---|
| 223 | self.multi = multi |
|---|
| 224 | if encoding is None: |
|---|
| 225 | encoding = sys.getdefaultencoding() |
|---|
| 226 | self.encoding = encoding |
|---|
| 227 | self.errors = errors |
|---|
| 228 | self.decode_keys = decode_keys |
|---|
| 229 | |
|---|
| 230 | def _decode_key(self, key): |
|---|
| 231 | if self.decode_keys: |
|---|
| 232 | try: |
|---|
| 233 | key = key.decode(self.encoding, self.errors) |
|---|
| 234 | except AttributeError: |
|---|
| 235 | pass |
|---|
| 236 | return key |
|---|
| 237 | |
|---|
| 238 | def _decode_value(self, value): |
|---|
| 239 | """ |
|---|
| 240 | Decode the specified value to unicode. Assumes value is a ``str`` or |
|---|
| 241 | `FieldStorage`` object. |
|---|
| 242 | |
|---|
| 243 | ``FieldStorage`` objects are specially handled. |
|---|
| 244 | """ |
|---|
| 245 | if isinstance(value, cgi.FieldStorage): |
|---|
| 246 | # decode FieldStorage's field name and filename |
|---|
| 247 | value = copy.copy(value) |
|---|
| 248 | if self.decode_keys: |
|---|
| 249 | value.name = value.name.decode(self.encoding, self.errors) |
|---|
| 250 | value.filename = value.filename.decode(self.encoding, self.errors) |
|---|
| 251 | else: |
|---|
| 252 | try: |
|---|
| 253 | value = value.decode(self.encoding, self.errors) |
|---|
| 254 | except AttributeError: |
|---|
| 255 | pass |
|---|
| 256 | return value |
|---|
| 257 | |
|---|
| 258 | def __getitem__(self, key): |
|---|
| 259 | return self._decode_value(self.multi.__getitem__(key)) |
|---|
| 260 | |
|---|
| 261 | def __setitem__(self, key, value): |
|---|
| 262 | self.multi.__setitem__(key, value) |
|---|
| 263 | |
|---|
| 264 | def add(self, key, value): |
|---|
| 265 | """ |
|---|
| 266 | Add the key and value, not overwriting any previous value. |
|---|
| 267 | """ |
|---|
| 268 | self.multi.add(key, value) |
|---|
| 269 | |
|---|
| 270 | def getall(self, key): |
|---|
| 271 | """ |
|---|
| 272 | Return a list of all values matching the key (may be an empty list) |
|---|
| 273 | """ |
|---|
| 274 | return [self._decode_value(v) for v in self.multi.getall(key)] |
|---|
| 275 | |
|---|
| 276 | def getone(self, key): |
|---|
| 277 | """ |
|---|
| 278 | Get one value matching the key, raising a KeyError if multiple |
|---|
| 279 | values were found. |
|---|
| 280 | """ |
|---|
| 281 | return self._decode_value(self.multi.getone(key)) |
|---|
| 282 | |
|---|
| 283 | def mixed(self): |
|---|
| 284 | """ |
|---|
| 285 | Returns a dictionary where the values are either single |
|---|
| 286 | values, or a list of values when a key/value appears more than |
|---|
| 287 | once in this dictionary. This is similar to the kind of |
|---|
| 288 | dictionary often used to represent the variables in a web |
|---|
| 289 | request. |
|---|
| 290 | """ |
|---|
| 291 | unicode_mixed = {} |
|---|
| 292 | for key, value in self.multi.mixed().iteritems(): |
|---|
| 293 | if isinstance(value, list): |
|---|
| 294 | value = [self._decode_value(value) for value in value] |
|---|
| 295 | else: |
|---|
| 296 | value = self._decode_value(value) |
|---|
| 297 | unicode_mixed[self._decode_key(key)] = value |
|---|
| 298 | return unicode_mixed |
|---|
| 299 | |
|---|
| 300 | def dict_of_lists(self): |
|---|
| 301 | """ |
|---|
| 302 | Returns a dictionary where each key is associated with a |
|---|
| 303 | list of values. |
|---|
| 304 | """ |
|---|
| 305 | unicode_dict = {} |
|---|
| 306 | for key, value in self.multi.dict_of_lists().iteritems(): |
|---|
| 307 | value = [self._decode_value(value) for value in value] |
|---|
| 308 | unicode_dict[self._decode_key(key)] = value |
|---|
| 309 | return unicode_dict |
|---|
| 310 | |
|---|
| 311 | def __delitem__(self, key): |
|---|
| 312 | self.multi.__delitem__(key) |
|---|
| 313 | |
|---|
| 314 | def __contains__(self, key): |
|---|
| 315 | return self.multi.__contains__(key) |
|---|
| 316 | |
|---|
| 317 | has_key = __contains__ |
|---|
| 318 | |
|---|
| 319 | def clear(self): |
|---|
| 320 | self.multi.clear() |
|---|
| 321 | |
|---|
| 322 | def copy(self): |
|---|
| 323 | return UnicodeMultiDict(self.multi.copy(), self.encoding, self.errors) |
|---|
| 324 | |
|---|
| 325 | def setdefault(self, key, default=None): |
|---|
| 326 | return self._decode_value(self.multi.setdefault(key, default)) |
|---|
| 327 | |
|---|
| 328 | def pop(self, key, *args): |
|---|
| 329 | return self._decode_value(self.multi.pop(key, *args)) |
|---|
| 330 | |
|---|
| 331 | def popitem(self): |
|---|
| 332 | k, v = self.multi.popitem() |
|---|
| 333 | return (self._decode_key(k), self._decode_value(v)) |
|---|
| 334 | |
|---|
| 335 | def __repr__(self): |
|---|
| 336 | items = ', '.join(['(%r, %r)' % v for v in self.items()]) |
|---|
| 337 | return '%s([%s])' % (self.__class__.__name__, items) |
|---|
| 338 | |
|---|
| 339 | def __len__(self): |
|---|
| 340 | return self.multi.__len__() |
|---|
| 341 | |
|---|
| 342 | ## |
|---|
| 343 | ## All the iteration: |
|---|
| 344 | ## |
|---|
| 345 | |
|---|
| 346 | def keys(self): |
|---|
| 347 | return [self._decode_key(k) for k in self.multi.iterkeys()] |
|---|
| 348 | |
|---|
| 349 | def iterkeys(self): |
|---|
| 350 | for k in self.multi.iterkeys(): |
|---|
| 351 | yield self._decode_key(k) |
|---|
| 352 | |
|---|
| 353 | __iter__ = iterkeys |
|---|
| 354 | |
|---|
| 355 | def items(self): |
|---|
| 356 | return [(self._decode_key(k), self._decode_value(v)) for \ |
|---|
| 357 | k, v in self.multi.iteritems()] |
|---|
| 358 | |
|---|
| 359 | def iteritems(self): |
|---|
| 360 | for k, v in self.multi.iteritems(): |
|---|
| 361 | yield (self._decode_key(k), self._decode_value(v)) |
|---|
| 362 | |
|---|
| 363 | def values(self): |
|---|
| 364 | return [self._decode_value(v) for v in self.multi.itervalues()] |
|---|
| 365 | |
|---|
| 366 | def itervalues(self): |
|---|
| 367 | for v in self.multi.itervalues(): |
|---|
| 368 | yield self._decode_value(v) |
|---|
| 369 | |
|---|
| 370 | __test__ = { |
|---|
| 371 | 'general': """ |
|---|
| 372 | >>> d = MultiDict(a=1, b=2) |
|---|
| 373 | >>> d['a'] |
|---|
| 374 | 1 |
|---|
| 375 | >>> d.getall('c') |
|---|
| 376 | [] |
|---|
| 377 | >>> d.add('a', 2) |
|---|
| 378 | >>> d['a'] |
|---|
| 379 | 1 |
|---|
| 380 | >>> d.getall('a') |
|---|
| 381 | [1, 2] |
|---|
| 382 | >>> d['b'] = 4 |
|---|
| 383 | >>> d.getall('b') |
|---|
| 384 | [4] |
|---|
| 385 | >>> d.keys() |
|---|
| 386 | ['a', 'a', 'b'] |
|---|
| 387 | >>> d.items() |
|---|
| 388 | [('a', 1), ('a', 2), ('b', 4)] |
|---|
| 389 | >>> d.mixed() |
|---|
| 390 | {'a': [1, 2], 'b': 4} |
|---|
| 391 | >>> MultiDict([('a', 'b')], c=2) |
|---|
| 392 | MultiDict([('a', 'b'), ('c', 2)]) |
|---|
| 393 | """} |
|---|
| 394 | |
|---|
| 395 | if __name__ == '__main__': |
|---|
| 396 | import doctest |
|---|
| 397 | doctest.testmod() |
|---|