Allow CompressedTextProperty to be None

README

@@ -7,5 +7,13 @@ generally useful. They include:
 - DerivedProperty, which allows you to automatically generate values
   - LowerCaseProperty, which stores the lower-cased value of another property
   - LengthProperty, which stores the length of another property
+- DerivedDateProperty is a date which is derived from a datetime property on the
+  same model.
+- ChoiceProperty efficiently handles properties which may only be assigned a
+  value from a limited set of choices
+- CompressedBlobProperty and CompressedTextProperty store data/text in a
+  compressed form
+- PacificDateTimeProperty is a DateTimeProperty whose value is always returned
+  in Pacific time.
 
 With aetycoon, you'll have all the properties you're ever likely to need.

__init__.py

@@ -1,8 +1,10 @@
 import copy
+import datetime
 import hashlib
 import logging
 import os
 import pickle
+import zlib
 from google.appengine.api import users
 from google.appengine.ext import db
 
@@ -490,3 +492,312 @@ def get_value_for_datastore(self, model_instance):
           "Domain '%s' attempting to allegally modify data for domain '%s'"
           % (os.environ['HTTP_HOST'], value))
     return value
+
+
+class ChoiceProperty(db.IntegerProperty):
+  """A property for efficiently storing choices made from a finite set.
+
+  This works by mapping each choice to an integer.  The choices must be hashable
+  (so that they can be efficiently mapped back to their corresponding index).
+
+  Example usage:
+
+  >>> class ChoiceModel(db.Model):
+  ...   a_choice = ChoiceProperty(enumerate(['red', 'green', 'blue']))
+  ...   b_choice = ChoiceProperty([(0,None), (1,'alpha'), (4,'beta')])
+
+  You interact with choice properties using the choice values:
+
+  >>> model = ChoiceModel(a_choice='green')
+  >>> model.a_choice
+  'green'
+  >>> model.b_choice == None
+  True
+  >>> model.b_choice = 'beta'
+  >>> model.b_choice
+  'beta'
+  >>> model.put() # doctest: +ELLIPSIS
+  datastore_types.Key.from_path(u'ChoiceModel', ...)
+
+  >>> model2 = ChoiceModel.all().get()
+  >>> model2.a_choice
+  'green'
+  >>> model.b_choice
+  'beta'
+
+  To get the int representation of a choice, you may use either access the
+  choice's corresponding attribute or use the c2i method:
+  >>> green = ChoiceModel.a_choice.GREEN
+  >>> none = ChoiceModel.b_choice.c2i(None)
+  >>> (green == 1) and (none == 0)
+  True
+
+  The int representation of a choice is needed to filter on a choice property:
+  >>> ChoiceModel.gql("WHERE a_choice = :1", green).count()
+  1
+  """
+  def __init__(self, choices, make_choice_attrs=True, *args, **kwargs):
+    """Constructor.
+
+    Args:
+      choices: A non-empty list of 2-tuples of the form (id, choice). id must be
+        the int to store in the database.  choice may be any hashable value.
+      make_choice_attrs: If True, the uppercase version of each string choice is
+        set as an attribute whose value is the choice's int representation.
+    """
+    super(ChoiceProperty, self).__init__(*args, **kwargs)
+    self.index_to_choice = dict(choices)
+    self.choice_to_index = dict((c,i) for i,c in self.index_to_choice.iteritems())
+    if make_choice_attrs:
+      for i,c in self.index_to_choice.iteritems():
+        if isinstance(c, basestring):
+          setattr(self, c.upper(), i)
+
+  def get_choices(self):
+    """Gets a list of values which may be assigned to this property."""
+    return self.choice_to_index.keys()
+
+  def c2i(self, choice):
+    """Converts a choice to its datastore representation."""
+    return self.choice_to_index[choice]
+
+  def __get__(self, model_instance, model_class):
+    if model_instance is None:
+      return self
+    index = super(ChoiceProperty, self).__get__(model_instance, model_class)
+    return self.index_to_choice[index]
+
+  def __set__(self, model_instance, value):
+    try:
+      index = self.c2i(value)
+    except KeyError:
+      raise db.BadValueError('Property %s must be one of the allowed choices: %s' %
+                          (self.name, self.get_choices()))
+    super(ChoiceProperty, self).__set__(model_instance, index)
+
+  def get_value_for_datastore(self, model_instance):
+    # just use the underlying value from the parent
+    return super(ChoiceProperty, self).__get__(model_instance, model_instance.__class__)
+
+  def make_value_from_datastore(self, value):
+    if value is None:
+      return None
+    return self.index_to_choice[value]
+
+
+class CompressedProperty(db.UnindexedProperty):
+  """A unindexed property that is stored in a compressed form.
+
+  CompressedTextProperty and CompressedBlobProperty derive from this class.
+  """
+  def __init__(self, level, *args, **kwargs):
+    """Constructor.
+
+    Args:
+    level: Controls the level of zlib's compression (between 1 and 9).
+    """
+    super(CompressedProperty, self).__init__(*args, **kwargs)
+    self.level = level
+
+  def get_value_for_datastore(self, model_instance):
+    value = self.value_to_str(model_instance)
+    if value is not None:
+      return db.Blob(zlib.compress(value, self.level))
+
+  def make_value_from_datastore(self, value):
+    if value is not None:
+      ds_value = zlib.decompress(value)
+      return self.str_to_value(ds_value)
+
+  # override value_to_str and str_to_value to implement a new CompressedProperty
+  def value_to_str(self, model_instance):
+    """Returns the value stored by this property encoded as a (byte) string,
+    or None if value is None.  This string will be stored in the datastore.
+    By default, returns the value unchanged."""
+    return self.__get__(model_instance, model_instance.__class__)
+
+  @staticmethod
+  def str_to_value(s):
+    """Reverse of value_to_str.  By default, returns s unchanged."""
+    return s
+
+class CompressedBlobProperty(CompressedProperty):
+  """A byte string that will be stored in a compressed form.
+
+  Example usage:
+
+  >>> class CompressedBlobModel(db.Model):
+  ...   v = CompressedBlobProperty()
+
+  You can create a CompressedBlobProperty and set its value with your raw byte
+  string (anything of type str).  You can also retrieve the (decompressed) value
+  by accessing the field.
+
+  >>> model = CompressedBlobModel(v='\x041\x9f\x11')
+  >>> model.v = 'green'
+  >>> model.v
+  'green'
+  >>> model.put() # doctest: +ELLIPSIS
+  datastore_types.Key.from_path(u'CompressedBlobModel', ...)
+
+  >>> model2 = CompressedBlobModel.all().get()
+  >>> model2.v
+  'green'
+
+  Compressed blobs are not indexed and therefore cannot be filtered on:
+
+  >>> CompressedBlobModel.gql("WHERE v = :1", 'green').count()
+  0
+  """
+  data_type = db.Blob
+
+  def __init__(self, level=6, *args, **kwargs):
+    super(CompressedBlobProperty, self).__init__(level, *args, **kwargs)
+
+class CompressedTextProperty(CompressedProperty):
+  """A string that will be stored in a compressed form (encoded as UTF-8).
+
+  Example usage:
+
+  >>> class CompressedTextModel(db.Model):
+  ...  v = CompressedTextProperty()
+
+  You can create a CompressedTextProperty and set its value with your string.
+  You can also retrieve the (decompressed) value by accessing the field.
+
+  >>> ustr = u'\u043f\u0440\u043e\u0440\u0438\u0446\u0430\u0442\u0435\u043b\u044c'
+  >>> model = CompressedTextModel(v=ustr)
+  >>> model.put() # doctest: +ELLIPSIS
+  datastore_types.Key.from_path(u'CompressedTextModel', ...)
+
+  >>> model2 = CompressedTextModel.all().get()
+  >>> model2.v == ustr
+  True
+
+  Compressed text is not indexed and therefore cannot be filtered on:
+
+  >>> CompressedTextModel.gql("WHERE v = :1", ustr).count()
+  0
+  """
+  data_type = db.Text
+
+  def __init__(self, level=6, *args, **kwargs):
+    super(CompressedTextProperty, self).__init__(level, *args, **kwargs)
+
+  def value_to_str(self, model_instance):
+    v = self.__get__(model_instance, model_instance.__class__)
+    if v is None:
+      return v
+    return v.encode('utf-8')
+
+  @staticmethod
+  def str_to_value(s):
+    if s is None:
+      return None
+    return s.decode('utf-8')
+
+class DerivedDateProperty(db.DateProperty):
+  """A date which is derived from a datetime property on the same model.
+
+  This is useful when you want to denormalize by storing a plain date which can
+  be queried for equality (while also preserving the more detailed time
+  information in the original datetime property).
+
+  This could also be accomplished with a normal DerivedProperty, though the
+  return type would be datetime rather than a plain date.
+  """
+  def __init__(self, dt_prop_name, *args, **kwargs):
+    super(DerivedDateProperty, self).__init__(*args, **kwargs)
+    self.dt_prop_name = dt_prop_name
+
+  def __get__(self, model_instance, model_class):
+    if model_instance is None:
+      return self
+    return getattr(model_instance, self.dt_prop_name).date()
+
+  def __set__(self, model_instance, value):
+    raise db.DerivedPropertyError("Cannot assign to a DerivedProperty")
+
+class UTC_tzinfo(datetime.tzinfo):
+    def utcoffset(self, dt):
+        return datetime.timedelta(0)
+
+    def dst(self, dt):
+        return datetime.timedelta(0)
+
+    def tzname(self, dt):
+        return "UTC"
+
+# source: http://code.google.com/appengine/docs/python/datastore/typesandpropertyclasses.html
+class PT_tzinfo(datetime.tzinfo):
+  """Implementation of the Pacific timezone."""
+  def utcoffset(self, dt):
+    return datetime.timedelta(hours=-8) + self.dst(dt)
+
+  def _FirstSunday(self, dt):
+    """First Sunday on or after dt."""
+    return dt + datetime.timedelta(days=(6-dt.weekday()))
+
+  def dst(self, dt):
+    # 2 am on the second Sunday in March
+    dst_start = self._FirstSunday(datetime.datetime(dt.year, 3, 8, 2))
+    # 1 am on the first Sunday in November
+    dst_end = self._FirstSunday(datetime.datetime(dt.year, 11, 1, 1))
+
+    if dst_start <= dt.replace(tzinfo=None) < dst_end:
+      return datetime.timedelta(hours=1)
+    else:
+      return datetime.timedelta(hours=0)
+
+    def tzname(self, dt):
+      if self.dst(dt) == datetime.timedelta(hours=0):
+        return "PST"
+      else:
+        return "PDT"
+
+TZ_UTC = UTC_tzinfo()
+TZ_PT = PT_tzinfo()
+
+def naive_dt_to_pt(naive_dt):
+    """
+    Takes a naive datetime object, assigns UTC timezone to it, and then
+    converts it to pacific time.  This is appropriate for datetime objects
+    which are retrieved from the datastore (they never include TZ info and
+    are ALWAYS in UTC, even if the object being put had other TZ info attached).
+    """
+    utc_dt = naive_dt.replace(tzinfo=TZ_UTC)
+    return utc_dt.astimezone(TZ_PT)
+
+class PacificDateTimeProperty(db.DateTimeProperty):
+  """A DateTimeProperty whose value is always returned in Pacific time.
+
+  Example usage:
+
+  >>> class TestModel(db.Model):
+  ...   dt = PacificDateTimeProperty()
+  ...   date = DerivedDateProperty('dt')
+
+  >>> dt = datetime.datetime(2010, 03, 25, 4, 0, 0)  # 4am UTC on 2010-Mar-25
+  >>> e = TestModel(dt=dt)
+  >>> e.put() # doctest: +ELLIPSIS
+  datastore_types.Key.from_path(u'TestModel', ...)
+
+  >>> print e.dt
+  2010-03-24 21:00:00-07:00
+
+  >>> print e.date
+  2010-03-24
+  """
+  def __get__(self, model_instance, model_class):
+    if model_instance is None:
+      return self
+    dt = super(PacificDateTimeProperty, self).__get__(model_instance, model_class)
+    if dt.tzinfo is None or dt.tzinfo.utcoffset(dt) is None:
+      return naive_dt_to_pt(dt)
+    else:
+      return dt  # don't convert it if it has already been converted
+
+  def make_value_from_datastore(self, value):
+    if value is None:
+      return None
+    return naive_dt_to_pt(value)