Skip to content

Instantly share code, notes, and snippets.

@pySilver
Last active November 26, 2019 22:56
Show Gist options
  • Save pySilver/a2b17f8cb42c18b5c4f2976f2a65b5e2 to your computer and use it in GitHub Desktop.
Save pySilver/a2b17f8cb42c18b5c4f2976f2a65b5e2 to your computer and use it in GitHub Desktop.
ModelChoiceBlock & MultipleModelChoiceBlock for Wagtail CMS
from typing import List, Optional
from django import forms
from django.db.models import Model, QuerySet
from django.utils.functional import cached_property
from wagtail.core.blocks import FieldBlock
from wagtail.core.utils import resolve_model_string
from dal_select2.widgets import ModelSelect2, ModelSelect2Multiple
class ModelChoiceBlock(FieldBlock):
class Meta:
icon = "placeholder"
def __init__(
self,
target_model=None,
queryset=None,
to_field_name=None,
limit_choices_to=None,
required=True,
help_text=None,
validators=(),
widget=None,
**kwargs
):
self._target_model = target_model
self._queryset = queryset
self._to_field_name = to_field_name
self._limit_choices_to = limit_choices_to
self._required = required
self._help_text = help_text
self._validators = validators
if widget is None:
widget = ModelSelect2
self._widget = widget
# keep a copy of all kwargs (including our normalised choices list) for deconstruct()
self._constructor_kwargs = kwargs.copy()
if target_model is not None:
self._constructor_kwargs["target_model"] = target_model
if queryset is not None:
self._constructor_kwargs["queryset"] = queryset
if to_field_name is not None:
self._constructor_kwargs["to_field_name"] = to_field_name
if limit_choices_to is not None:
self._constructor_kwargs["limit_choices_to"] = limit_choices_to
if required is not True:
self._constructor_kwargs["required"] = required
if help_text is not None:
self._constructor_kwargs["help_text"] = help_text
super().__init__(**kwargs)
@cached_property
def queryset(self) -> QuerySet:
queryset = None
if self._target_model is not None:
queryset = resolve_model_string(self._target_model).objects.all()
if self._queryset is not None:
queryset = self._queryset
if callable(self._queryset):
queryset = self._queryset()
if queryset is None:
raise ValueError("target_model or queryset should be provided.")
return queryset
@cached_property
def field(self) -> "forms.ModelChoiceField":
return forms.ModelChoiceField(
queryset=self.queryset,
limit_choices_to=self._limit_choices_to,
to_field_name=self._to_field_name,
widget=self._widget,
required=self._required,
validators=self._validators,
help_text=self._help_text,
)
def clean(self, value):
if isinstance(value, self.queryset.model):
key = self._to_field_name or "pk"
value = getattr(value, key)
return super().clean(value)
def to_python(self, value) -> Optional[Model]:
# the incoming serialised value should be None or a specific field value
if value is None:
return value
else:
try:
key = self._to_field_name or "pk"
return self.queryset.model.objects.get(**{key: value})
except self.queryset.model.DoesNotExist:
return None
def bulk_to_python(self, values) -> List[Model]:
"""Return the model instances for the given list of primary keys.
The instances must be returned in the same order as the values and keep None values.
"""
key = self._to_field_name or "pk"
objects = self.queryset.model.objects.in_bulk(values, field_name=key)
return [
objects.get(key) for key in values
] # Keeps the ordering the same as in values.
def get_prep_value(self, value):
# the native value (a model instance or None) should serialise to a specific field value or None
if value is None:
return None
else:
key = self._to_field_name or "pk"
return getattr(value, key)
def value_from_form(self, value):
# ModelChoiceField sometimes returns scalar value, and sometimes an instance; we want the instance
if value is None or isinstance(value, self.queryset.model):
return value
else:
try:
key = self._to_field_name or "pk"
return self.queryset.model.objects.get(**{key: value})
except self.queryset.model.DoesNotExist:
return None
def deconstruct(self):
return ("wagtail.core.blocks.ModelChoiceBlock", [], self._constructor_kwargs)
class ModelMultipleChoiceBlock(ModelChoiceBlock):
class Meta:
icon = "placeholder"
def __init__(
self,
target_model=None,
queryset=None,
to_field_name=None,
limit_choices_to=None,
required=True,
help_text=None,
validators=(),
widget=None,
**kwargs
):
if widget is None:
widget = ModelSelect2Multiple
super().__init__(
target_model,
queryset,
to_field_name,
limit_choices_to,
required,
help_text,
validators,
widget,
**kwargs
)
@cached_property
def field(self) -> "forms.ModelMultipleChoiceField":
return forms.ModelMultipleChoiceField(
queryset=self.queryset,
limit_choices_to=self._limit_choices_to,
to_field_name=self._to_field_name,
widget=self._widget,
required=self._required,
validators=self._validators,
help_text=self._help_text,
)
def clean(self, value):
if isinstance(value, QuerySet):
key = self._to_field_name or "pk"
value = [getattr(obj, key) for obj in value]
return super().clean(value)
def get_prep_value(self, value):
# the native value (a queryset or None) should serialise to a list of specific field values or None
if value is None:
return None
else:
key = self._to_field_name or "pk"
return [getattr(obj, key) for obj in value]
def to_python(self, value):
# the incoming serialised value should be None or a list of specific field values
if value is None:
return value
else:
try:
key = self._to_field_name or "pk"
return self.queryset.model.objects.filter(**{"%s__in" % key: value})
except self.queryset.model.DoesNotExist:
return None
def value_from_form(self, value):
# ModelChoiceField sometimes returns list of scalar values, and sometimes a queryset; we want the queryset
if value is None or isinstance(value, QuerySet):
return value
else:
try:
key = self._to_field_name or "pk"
return self.queryset.model.objects.filter(**{"%s__in" % key: value})
except self.queryset.model.DoesNotExist:
return None
def deconstruct(self):
return (
"wagtail.core.blocks.ModelMultipleChoiceBlock",
[],
self._constructor_kwargs,
)
@pySilver
Copy link
Author

Usage:

brand = ModelChoiceBlock(
        label=_("Brand"),
        target_model="brands.Brand",  # or model 
        required=False,
        to_field_name="slug",
    )

or

brand = ModelChoiceBlock(
        label=_("Brand"),
        queryset=Brand.objects.all(),  # or callable that returns queryset
        required=False,
        to_field_name="slug",
    )

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment