forum/models/utils.py

   1 from django.db import models
   2 from django.core.cache import cache
   3 from django.conf import settings
   4 from django.utils.encoding import force_unicode
   5
   6 try:
   7     from cPickle import loads, dumps
   8 except ImportError:
   9     from pickle import loads, dumps
  10
  11 from copy import deepcopy
  12 from base64 import b64encode, b64decode
  13 from zlib import compress, decompress
  14 import re
  15
  16 from base import BaseModel
  17
  18 MAX_MARKABLE_STRING_LENGTH = 100
  19
  20 class PickledObject(unicode):
  21     pass
  22
  23 def dbsafe_encode(value, compress_object=True):
  24     if not compress_object:
  25         value = b64encode(dumps(deepcopy(value)))
  26     else:
  27         value = b64encode(compress(dumps(deepcopy(value))))
  28     return PickledObject(value)
  29
  30 def dbsafe_decode(value, compress_object=True):
  31     if not compress_object:
  32         value = loads(b64decode(value))
  33     else:
  34         value = loads(decompress(b64decode(value)))
  35     return value
  36
  37 class PickledObjectField(models.Field):
  38     __metaclass__ = models.SubfieldBase
  39
  40     marker_re = re.compile(r'^T\[(?P<type>\w+)\](?P<value>.*)$', re.DOTALL)
  41     markable_types = dict((t.__name__, t) for t in (str, int, unicode))
  42
  43     def __init__(self, *args, **kwargs):
  44         self.compress = kwargs.pop('compress', True)
  45         self.protocol = kwargs.pop('protocol', 2)
  46         kwargs.setdefault('null', True)
  47         kwargs.setdefault('editable', False)
  48         super(PickledObjectField, self).__init__(*args, **kwargs)
  49
  50     def generate_type_marked_value(self, value):
  51         return PickledObject(u"T[%s]%s" % (type(value).__name__, value))
  52
  53     def read_marked_value(self, value):
  54         m = self.marker_re.match(value)
  55
  56         if m:
  57             marker = m.group('type')
  58             value = m.group('value')
  59             if marker in self.markable_types:
  60                 value = self.markable_types[marker](value)
  61
  62         return value
  63
  64     def get_default(self):
  65         if self.has_default():
  66             if callable(self.default):
  67                 return self.default()
  68             return self.default
  69
  70         return super(PickledObjectField, self).get_default()
  71
  72     def to_python(self, value):
  73         if value is not None:
  74             try:
  75                 if value.startswith("T["):
  76                     value = self.read_marked_value(value)
  77                 else:
  78                     value = dbsafe_decode(value, self.compress)
  79             except:
  80                 if isinstance(value, PickledObject):
  81                     raise
  82         return value
  83
  84     def get_db_prep_value(self, value, connection, prepared=False):
  85         if value is not None and not isinstance(value, PickledObject):
  86             if type(value).__name__ in self.markable_types and not (isinstance(value, basestring) and len(value
  87                                                                                                           ) > MAX_MARKABLE_STRING_LENGTH):
  88                 value = unicode(self.generate_type_marked_value(value))
  89             else:
  90                 value = unicode(dbsafe_encode(value, self.compress))
  91         return value
  92
  93     def value_to_string(self, obj):
  94         value = self._get_val_from_obj(obj)
  95         return self.get_db_prep_value(value)
  96
  97     def get_internal_type(self):
  98         return 'TextField'
  99
 100     def get_db_prep_lookup(self, lookup_type, value):
 101         if lookup_type not in ['exact', 'in', 'isnull']:
 102             raise TypeError('Lookup type %s is not supported.' % lookup_type)
 103         return super(PickledObjectField, self).get_db_prep_lookup(lookup_type, value)
 104
 105
 106 class KeyValue(BaseModel):
 107     key = models.CharField(max_length=255, unique=True)
 108     value = PickledObjectField()
 109
 110     class Meta:
 111         app_label = 'forum'
 112
 113     def cache_key(self):
 114         return self._generate_cache_key(self.key)
 115
 116     @classmethod
 117     def infer_cache_key(cls, querydict):
 118         try:
 119             key = [v for (k, v) in querydict.items() if k in ('key', 'key__exact')][0]
 120
 121             return cls._generate_cache_key(key)
 122         except:
 123             return None
 124
 125     @classmethod
 126     def value_to_list_on_cache_query(cls):
 127         return 'key'
 128