Code

Ticket #14969: modify_models.py

File modify_models.py, 3.9 KB (added by marinho, 4 years ago)
Line 
1"""
2Examples
3========
4
5my_app/models.py
6----------------
7
8    from django.db import models
9
10    class CustomerType(models.Model):
11        name = models.CharField(max_length=50)
12
13        def __unicode__(self):
14            return self.name
15
16    class Customer(models.Model):
17        name = models.CharField(max_length=50)
18        type = models.ForeignKey('CustomerType')
19        is_active = models.BooleanField(default=True, blank=True)
20        employer = models.CharField(max_length=100)
21
22        def __unicode__(self):
23            return self.name
24
25another_app/models.py
26---------------------
27
28    from django.db import models
29    from django.contrib.auth.models import User
30
31    from djangoplus.modify_models import ModifiedModel
32
33    class City(models.Model):
34        name = models.CharField(max_length=50)
35
36        def __unicode__(self):
37            return self.name
38
39    class HelperCustomerType(ModifiedModel):
40        class Meta:
41            model = 'my_app.CustomerType'
42
43        description = models.TextField()
44
45    class HelperCustomer(ModifiedModel):
46        class Meta:
47            model = 'my_app.Customer'
48            exclude = ('employer',)
49
50        type = models.CharField(max_length=50)
51        address = models.CharField(max_length=100)
52        city = models.ForeignKey(City)
53
54        def __unicode__(self):
55            return '%s - %s'%(self.pk, self.name)
56
57    class HelperUser(ModifiedModel):
58        class Meta:
59            model = User
60
61        website = models.URLField(blank=True, verify_exists=False)
62
63"""
64import types
65
66from django.db import models
67from django.core.exceptions import ImproperlyConfigured
68from django.db.models import get_model
69from django.db.models.fields import FieldDoesNotExist
70
71class ModifiedModelMetaclass(type):
72    def __new__(cls, name, bases, attrs):
73        new_class = super(ModifiedModelMetaclass, cls).__new__(cls, name, bases, attrs)
74
75        if name == 'ModifiedModel' and bases[0] == object:
76            return new_class
77
78        try:
79            meta = attrs['Meta']()
80        except KeyError:
81            raise ImproperlyConfigured("Helper class %s hasn't a Meta subclass!" % name)
82
83        # Find model class for this helper
84        if isinstance(getattr(meta, 'model', None), basestring):
85            model_class = get_model(*meta.model.split('.'))
86        elif issubclass(getattr(meta, 'model', None), models.Model):
87            model_class = meta.model
88        else:
89            raise ImproperlyConfigured("Model informed by Meta subclass of %s is improperly!" % name)
90
91        def remove_field(f_name):
92            # Removes the field form local fields list
93            model_class._meta.local_fields = [f for f in model_class._meta.local_fields
94                    if f.name != f_name]
95
96            # Removes the field setter if exists
97            if hasattr(model_class, f_name):
98                delattr(model_class, f_name)
99
100        # Removes fields setted in attribute 'exclude'
101        if isinstance(getattr(meta, 'exclude', None), (list,tuple)):
102            for f_name in meta.exclude:
103                remove_field(f_name)
104
105        # Calls 'contribute_to_class' from field to sender class
106        for f_name, field in attrs.items():
107            if isinstance(field, models.Field):
108                # Removes the field if it already exists
109                remove_field(f_name)
110
111                # Appends the new field to model class
112                field.contribute_to_class(model_class, f_name)
113
114        # Attaches methods
115        for m_name, func in attrs.items():
116            if callable(func) and type(func) == types.FunctionType:
117                setattr(model_class, m_name, func)
118
119        new_class._meta = meta
120
121        return new_class
122
123class ModifiedModel(object):
124    """
125    Make your inheritance from this class and set a Meta subclass with attribute
126    'model' with the model class you want to modify: add/replace/exclude fields
127    and/or add/replace methods.
128    """
129    __metaclass__ = ModifiedModelMetaclass
130