diff --git a/categories/__init__.py b/categories/__init__.py index e69de29..5aab847 100644 --- a/categories/__init__.py +++ b/categories/__init__.py @@ -0,0 +1,18 @@ +import fields +import models + +from django.db.models import FieldDoesNotExist + +def register_m2m(model, field_name, extra_params={}): + return _register(model, field_name, extra_params, fields.CategoryM2MField) + + +def register_fk(model, field_name, extra_params={}): + return _register(model, field_name, extra_params, fields.CategoryFKField) + +def _register(model, field_name='category', extra_params={}, field=fields.CategoryFKField): + opts = model._meta + try: + opts.get_field(field_name) + except FieldDoesNotExist: + field(**extra_params).contribute_to_class(model, field_name) \ No newline at end of file diff --git a/categories/fields.py b/categories/fields.py new file mode 100644 index 0000000..fde87ca --- /dev/null +++ b/categories/fields.py @@ -0,0 +1,11 @@ +from django.db.models import ForeignKey, ManyToManyField +from categories.models import Category + +class CategoryM2MField(ManyToManyField): + def __init__(self, **kwargs): + super(CategoryM2MField, self).__init__(to=Category, **kwargs) + + +class CategoryFKField(ForeignKey): + def __init__(self, **kwargs): + super(CategoryFKField, self).__init__(to=Category, **kwargs) \ No newline at end of file