Django snippets

import datetime
import warnings
import pytz
from django.conf import settings
from django.db import models
from django.dispatch import dispatcher
from django.db.models import signals

try:
    pytz_exc = (pytz.UnknownTimeZoneError,)
except AttributeError:
    # older versions of pytz
    pytz_exc = (IOError, KeyError)

TIMEZONES = tuple((t, t) for t in pytz.all_timezones)


class UTCDateTimeField(models.DateTimeField):
    
    def __init__(self, *args, **kw):
        self.default_tz = kw.get('default_tz')
        if 'default_tz' in kw:
            del(kw['default_tz'])
        super(UTCDateTimeField, self).__init__(*args, **kw)
        self._name = None
        
    def get_internal_type(self):
        return 'DateTimeField'

    def _model_init(self, signal, sender, instance):
        setattr(instance, "_%s" % self._name, getattr(instance, self._name))
        setattr(instance, "_%s_utc" % self._name, getattr(instance, "%s_utc" % self._name))
        setattr(instance, "_%s_tz" % self._name, getattr(instance, "%s_tz" % self._name))
        
    def _model_pre_save(self, signal, sender, instance):
        dt = getattr(instance, self._name)
        utc = getattr(instance, "%s_utc" % self._name)
        tz = getattr(instance, "%s_tz" % self._name)
        if not tz:
            # try to get the default
            if isinstance(self.default_tz, basestring):
                tz = getattr(instance, self.default_tz)
                if isinstance(tz, basestring):
                    # use it as a timezone
                    pass
                elif callable(tz):
                    self.default_tz = tz
                    tz = None
                else:
                    warnings.warn("UTCDateTimeField cannot use default_tz %s in model %s" % (self.default_tz, instance))
                    self.default_tz = lambda: getattr(settings, 'TIME_ZONE', 'UTC')
                    tz = None
            if callable(self.default_tz):
                tz = self.default_tz()
                if not tz or not isinstance(tz, basestring):
                    warnings.warn("UTCDateTimeField cannot use default timezone %s, using Django default" % tz)
                    tz = getattr(settings, 'TIME_ZONE', 'UTC')
            if not tz:
                tz = getattr(settings, 'TIME_ZONE', 'UTC')
            setattr(instance, "%s_tz" % self._name, tz)
        _dt = getattr(instance, "_%s" % self._name)
        _utc = getattr(instance, "_%s_utc" % self._name)
        _tz = getattr(instance, "_%s_tz" % self._name)
        if isinstance(tz, unicode):
            tz = tz.encode('utf8')
        try:
            timezone = pytz.timezone(tz)
        except pytz_exc, e:
            warnings.warn("Error in timezone: %s" % e)
            setattr(instance, "%s_tz" % self._name, 'UTC')
            timezone = pytz.UTC
        if not dt and not utc:
            # do nothing
            return
        elif dt and dt != _dt:
            utc = UTCDateTimeField.utc_from_localtime(dt, timezone)
            setattr(instance, "%s_utc" % self._name, utc)
        elif utc and (utc != _utc or tz != _tz):
            dt = UTCDateTimeField.localtime_from_utc(utc, timezone)
            setattr(instance, "_%s" % self._name, dt)
        elif not dt:
            dt = UTCDateTimeField.localtime_from_utc(utc, timezone)
            setattr(instance, "_%s" % self._name, dt)
        elif not utc:
            utc = UTCDateTimeField.utc_from_localtime(dt, timezone)
            setattr(instance, "%s_utc" % self._name, utc)
        if dt.tzinfo:
            dt = dt.replace(tzinfo=None)
        if utc.tzinfo:
            utc = utc.replace(tzinfo=None)
        setattr(instance, "_%s" % self._name, dt)
        setattr(instance, "_%s_utc" % self._name, utc)
        setattr(instance, "_%s_tz" % self._name, tz)
        setattr(instance, "%s" % self._name, dt)
        setattr(instance, "%s_utc" % self._name, utc)
        
    def _model_post_save(self, signal, sender, instance):
        setattr(instance, "%s" % self._name, getattr(instance, "_%s" % self._name))
        setattr(instance, "%s_utc" % self._name, getattr(instance, "_%s_utc" % self._name))
        
    
    def contribute_to_class(self, cls, name):
        self._name = name

        super(UTCDateTimeField, self).contribute_to_class(cls, name)
        self.utc_field = models.DateTimeField(blank=True)
        self.creation_counter = self.utc_field.creation_counter + 1
        models.DateTimeField.contribute_to_class(self.utc_field, cls, '%s_utc' % name)
        self.tz_field = models.CharField(maxlength=38, blank=True, choices=TIMEZONES, default=getattr(settings, 'TIME_ZONE', 'UTC'))
        self.creation_counter = self.tz_field.creation_counter + 1
        models.CharField.contribute_to_class(self.tz_field, cls, '%s_tz' % name)
        
        # add the properties for offset-aware datetimes
        def get_timezone(s):
            tz = getattr(s, '%s_tz' % name)
            if isinstance(tz, unicode):
                tz = tz.encode('utf8')
            return pytz.timezone(tz)
        setattr(cls, "%s_timezone" % name, property(get_timezone))
        
        def get_dt_offset_aware(s):
            dt = getattr(s, "%s_utc_offset_aware" % name)
            tz = getattr(s, "%s_timezone" % name)
            return dt.astimezone(tz)
        setattr(cls, "%s_offset_aware" % name, property(get_dt_offset_aware))
        
        def get_utc_offset_aware(s):
            return getattr(s, '%s_utc' % name).replace(tzinfo=pytz.utc)
        setattr(cls, "%s_utc_offset_aware" % name, property(get_utc_offset_aware))
        
        dispatcher.connect(self._model_init, signal=signals.post_init, sender=cls)
        dispatcher.connect(self._model_pre_save, signal=signals.pre_save, sender=cls)
        dispatcher.connect(self._model_post_save, signal=signals.post_save, sender=cls)

    @staticmethod
    def localtime_from_utc(utc, tz):
        dt = utc.replace(tzinfo=pytz.utc)
        return dt.astimezone(tz)
        
    @staticmethod
    def utc_from_localtime(dt, tz):
        dt = dt.replace(tzinfo=tz)
        _dt = tz.normalize(dt)
        if dt.tzinfo != _dt.tzinfo:
            # Houston, we have a problem...
            # find out which one has a dst offset
            if _dt.tzinfo.dst(_dt):
                _dt -= _dt.tzinfo.dst(_dt)
            else:
                _dt += dt.tzinfo.dst(dt)
        return _dt.astimezone(pytz.utc)