Skip to content
This repository
branch: master
file 153 lines (132 sloc) 5.264 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152
import uuid

from django import forms
from django.db.models import Field, SubfieldBase
try:
    from django.utils.encoding import smart_unicode
except ImportError:
    from django.utils.encoding import smart_text as smart_unicode

try:
    # psycopg2 needs us to register the uuid type
    import psycopg2.extras
    psycopg2.extras.register_uuid()
except (ImportError, AttributeError):
    pass


class StringUUID(uuid.UUID):
    def __init__(self, *args, **kwargs):
        # get around UUID's immutable setter
        object.__setattr__(self, 'hyphenate', kwargs.pop('hyphenate', False))

        super(StringUUID, self).__init__(*args, **kwargs)

    def __unicode__(self):
        return unicode(str(self))

    def __str__(self):
        if self.hyphenate:
            return super(StringUUID, self).__str__()

        return self.hex

    def __len__(self):
        return len(self.__unicode__())


class UUIDField(Field):
    """
A field which stores a UUID value in hex format. This may also have
the Boolean attribute 'auto' which will set the value on initial save to a
new UUID value (calculated using the UUID1 method). Note that while all
UUIDs are expected to be unique we enforce this with a DB constraint.
"""
    # TODO: support binary storage types
    __metaclass__ = SubfieldBase

    def __init__(self, version=4, node=None, clock_seq=None,
            namespace=None, name=None, auto=False, hyphenate=False, *args, **kwargs):
        assert version in (1, 3, 4, 5), "UUID version %s is not supported." % version
        self.auto = auto
        self.version = version
        self.hyphenate = hyphenate
        # We store UUIDs in hex format, which is fixed at 32 characters.
        kwargs['max_length'] = 32
        if auto:
            # Do not let the user edit UUIDs if they are auto-assigned.
            kwargs['editable'] = False
            kwargs['blank'] = True
            kwargs['unique'] = True
        if version == 1:
            self.node, self.clock_seq = node, clock_seq
        elif version in (3, 5):
            self.namespace, self.name = namespace, name
        super(UUIDField, self).__init__(*args, **kwargs)

    def _create_uuid(self):
        if self.version == 1:
            args = (self.node, self.clock_seq)
        elif self.version in (3, 5):
            error = None
            if self.name is None:
                error_attr = 'name'
            elif self.namespace is None:
                error_attr = 'namespace'
            if error is not None:
                raise ValueError("The %s parameter of %s needs to be set." %
                                 (error_attr, self))
            if not isinstance(self.namespace, uuid.UUID):
                raise ValueError("The name parameter of %s must be an "
                                 "UUID instance." % self)
            args = (self.namespace, self.name)
        else:
            args = ()
        return getattr(uuid, 'uuid%s' % self.version)(*args)

    def db_type(self, connection=None):
        """
Return the special uuid data type on Postgres databases.
"""
        if connection and 'postgres' in connection.vendor:
            return 'uuid'
        return 'char(%s)' % self.max_length

    def pre_save(self, model_instance, add):
        """
This is used to ensure that we auto-set values if required.
See CharField.pre_save
"""
        value = getattr(model_instance, self.attname, None)
        if self.auto and add and not value:
            # Assign a new value for this attribute if required.
            uuid = self._create_uuid()
            setattr(model_instance, self.attname, uuid)
            value = uuid.hex
        return value

    def get_db_prep_value(self, value, connection, prepared=False):
        """
Casts uuid.UUID values into the format expected by the back end
"""
        if isinstance(value, uuid.UUID):
            value = str(value)
        if isinstance(value, str):
            if '-' in value:
                return value.replace('-', '')
        return value

    def value_to_string(self, obj):
        val = self._get_val_from_obj(obj)
        if val is None:
            data = ''
        else:
            data = unicode(val)
        return data

    def to_python(self, value):
        """
Returns a ``StringUUID`` instance from the value returned by the
database. This doesn't use uuid.UUID directly for backwards
compatibility, as ``StringUUID`` implements ``__unicode__`` with
``uuid.UUID.hex()``.
"""
        if not value:
            return None
        # attempt to parse a UUID including cases in which value is a UUID
        # instance already to be able to get our StringUUID in.
        return StringUUID(smart_unicode(value), hyphenate=self.hyphenate)

    def formfield(self, **kwargs):
        defaults = {
            'form_class': forms.CharField,
            'max_length': self.max_length,
        }
        defaults.update(kwargs)
        return super(UUIDField, self).formfield(**defaults)

try:
    from south.modelsinspector import add_introspection_rules
    add_introspection_rules([], [r"^uuidfield\.fields\.UUIDField"])
except ImportError:
    pass
Something went wrong with that request. Please try again.