awx/awx/main/fields.py

40 lines
1.5 KiB
Python

# Copyright (c) 2013 AnsibleWorks, Inc.
# All Rights Reserved.
# Django
from django.db import models
from django.db.models.fields.related import SingleRelatedObjectDescriptor
# South
from south.modelsinspector import add_introspection_rules
__all__ = ['AutoOneToOneField']
# Based on AutoOneToOneField from django-annoying:
# https://bitbucket.org/offline/django-annoying/src/a0de8b294db3/annoying/fields.py
class AutoSingleRelatedObjectDescriptor(SingleRelatedObjectDescriptor):
"""Descriptor for access to the object from its related class."""
def __get__(self, instance, instance_type=None):
try:
return super(AutoSingleRelatedObjectDescriptor,
self).__get__(instance, instance_type)
except self.related.model.DoesNotExist:
obj = self.related.model(**{self.related.field.name: instance})
if self.related.field.rel.parent_link:
for f in instance._meta.local_fields:
setattr(obj, f.name, getattr(instance, f.name))
obj.save()
return obj
class AutoOneToOneField(models.OneToOneField):
"""OneToOneField that creates related object if it doesn't exist."""
def contribute_to_related_class(self, cls, related):
setattr(cls, related.get_accessor_name(),
AutoSingleRelatedObjectDescriptor(related))
add_introspection_rules([([AutoOneToOneField], [], {})],
[r'^awx\.main\.fields\.AutoOneToOneField'])