awx/awx/main/fields.py

35 lines
1.4 KiB
Python

# Copyright (c) 2015 Ansible, Inc.
# All Rights Reserved.
# Django
from django.db import models
from django.db.models.fields.related import SingleRelatedObjectDescriptor
__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.related_model.DoesNotExist:
obj = self.related.related_model(**{self.related.field.name: instance})
if self.related.field.rel.parent_link:
raise NotImplementedError('not supported with polymorphic!')
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))