diff --git a/src/c3nav/editor/wrappers.py b/src/c3nav/editor/wrappers.py index c6bdb478..00401e33 100644 --- a/src/c3nav/editor/wrappers.py +++ b/src/c3nav/editor/wrappers.py @@ -414,6 +414,11 @@ class BaseQueryWrapper(BaseWrapper): def __len__(self): return len(self.get_queryset()) + def create(self, *args, **kwargs): + obj = self.model(*args, **kwargs) + obj.save() + return obj + class ManagerWrapper(BaseQueryWrapper): def get_queryset(self): @@ -437,6 +442,10 @@ class RelatedManagerWrapper(ManagerWrapper): pass return self.get_queryset().all() + def create(self, *args, **kwargs): + kwargs[self._obj.field.name] = self.instance + super().create(*args, **kwargs) + class ManyRelatedManagerWrapper(RelatedManagerWrapper): def _check_through(self): @@ -475,6 +484,9 @@ class ManyRelatedManagerWrapper(RelatedManagerWrapper): def all(self): return self.model.objects.filter(**self._obj.core_filters) + def create(self, *args, **kwargs): + raise NotImplementedError + class QuerySetWrapper(BaseQueryWrapper): @property