X-Git-Url: https://git.openstreetmap.org./osqa.git/blobdiff_plain/1a949f7c97dc2f34c135f5cdf088df2927d3d652..dec306155823082d0e25e6abb10edc0cd269f0dd:/forum/models/utils.py

diff --git a/forum/models/utils.py b/forum/models/utils.py
index 984a55c..5ac8900 100644
--- a/forum/models/utils.py
+++ b/forum/models/utils.py
@@ -1,44 +1,84 @@
 from django.db import models
 from django.core.cache import cache
 from django.conf import settings
-#try:
-#	import cPickle as pickle
-#except ImportError:
-import pickle
-import django.dispatch
+from django.utils.encoding import force_unicode
+
+try:
+    from cPickle import loads, dumps
+except ImportError:
+    from pickle import loads, dumps
+
+from copy import deepcopy
+from base64 import b64encode, b64decode
+from zlib import compress, decompress
+
 
 class PickledObject(str):
-	pass
+    pass
+
+def dbsafe_encode(value, compress_object=True):
+    if not compress_object:
+        value = b64encode(dumps(deepcopy(value)))
+    else:
+        value = b64encode(compress(dumps(deepcopy(value))))
+    return PickledObject(value)
+
+def dbsafe_decode(value, compress_object=True):
+    if not compress_object:
+        value = loads(b64decode(value))
+    else:
+        value = loads(decompress(b64decode(value)))
+    return value
 
 class PickledObjectField(models.Field):
     __metaclass__ = models.SubfieldBase
 
+    def __init__(self, *args, **kwargs):
+        self.compress = kwargs.pop('compress', True)
+        self.protocol = kwargs.pop('protocol', 2)
+        kwargs.setdefault('null', True)
+        kwargs.setdefault('editable', False)
+        super(PickledObjectField, self).__init__(*args, **kwargs)
+
+    def get_default(self):
+        if self.has_default():
+            if callable(self.default):
+                return self.default()
+            return self.default
+
+        return super(PickledObjectField, self).get_default()
+
     def to_python(self, value):
-        if isinstance (value, PickledObject):
-            return value
-
-        try:
-            return pickle.loads(value.encode('utf-8'))
-        except:
-            return value
-    
-    def get_db_prep_save(self, value):
         if value is not None:
-            if isinstance(value, PickledObject):
-                return str(value)
-            else:
-			    value = pickle.dumps(value)
+            try:
+                value = dbsafe_decode(value, self.compress)
+            except:
+                if isinstance(value, PickledObject):
+                    raise
+        return value
 
+    def get_db_prep_value(self, value):
+        if value is not None and not isinstance(value, PickledObject):
+            value = force_unicode(dbsafe_encode(value, self.compress))
         return value
 
+    def value_to_string(self, obj):
+        value = self._get_val_from_obj(obj)
+        return self.get_db_prep_value(value)
+
     def get_internal_type(self):
         return 'TextField'
 
+    def get_db_prep_lookup(self, lookup_type, value):
+        if lookup_type not in ['exact', 'in', 'isnull']:
+            raise TypeError('Lookup type %s is not supported.' % lookup_type)
+        return super(PickledObjectField, self).get_db_prep_lookup(lookup_type, value)
+
 
 class KeyValueManager(models.Manager):
 
     def create_cache_key(self, key):
-        return "%s:key_value:%s" % (settings.APP_URL, key)
+        return "%s:keyvalue:%s" % (settings.APP_URL, key)
 
     def save_to_cache(self, instance):
         cache.set(self.create_cache_key(instance.key), instance, 2592000)