Cacher is the code snippet organizer for pro developers

# coding=utf-8
import six
import mongoengine as mng
import datetime
import pytz


class AutoDateTimeField(mng.DateTimeField):
    '''Auto transform to datetime(from raw datetime string)'''
    DEFAULT_FORMAT = "%Y-%m-%d %H:%M:%S"

    def __init__(self, *args, **kwargs):
        self.format = kwargs.pop('format', self.DEFAULT_FORMAT)
        self.store_type = kwargs.pop('store_type', datetime.datetime)
        super(AutoDateTimeField, self).__init__(*args, **kwargs)

    def to_datetime(self, value, **kwargs):
        if isinstance(value, six.string_types):
            return datetime.datetime.strptime(value, self.format)
        return value

    def to_mongo(self, value, **kwargs):
        '''for storage or lookup from storage'''
        value = self.to_datetime(value, **kwargs)
        if self.store_type is datetime.datetime:
            return value
        return value.strftime(self.format)

    def to_python(self, value):
        '''string and datetime types are all allowed in storage'''
        if value is None or isinstance(value, datetime.datetime):
            return value

        return datetime.datetime.strptime(value, self.format)


class AwareDateTimeField(AutoDateTimeField):
    def __init__(self, *args, **kwargs):
        self.timezone = pytz.timezone(kwargs.pop('timezone', 'Asia/Shanghai'))
        super(AwareDateTimeField, self).__init__(*args, **kwargs)

    def to_datetime(self, value, **kwargs):
        value = AutoDateTimeField.to_datetime(self, value, **kwargs)
        if value.tzinfo:
            return value.astimezone(pytz.utc)
        else:   # supposed to be naive local time
            return self.timezone.localize(value).astimezone(pytz.utc)

    def to_mongo(self, value, **kwargs):
        return self.to_datetime(value, **kwargs)

    def to_python(self, value):
        value = AutoDateTimeField.to_python(self, value)
        if value.tzinfo:
            return value.astimezone(self.timezone)
        else:
            return self.timezone.localize(value)

    # for motorengine
    to_son = to_mongo
    from_son = to_python


class RectifiableField(mng.base.BaseField):
    def __init__(self, *args, **kwargs):
        self.corrections = kwargs.pop('corrections')
        self.db_field_type = kwargs.pop('db_field_type')
        self.canonical_field_type = kwargs.pop('canonical_field_type')
        super(RectifiableField, self).__init__(*args, **kwargs)

    def to_python(self, value):
        return self.corrections.get(value)

    def to_mongo(self, value, **kwargs):
        for k, v in self.corrections.items():
            print value, v, k
            print type(value), type(v), type(k)
            if value == v:
                return k
        return value

    def prepare_query_value(self, op, value):
        return self.to_mongo(value)