ModelChoiceBlock & MultipleModelChoiceBlock for Wagtail CMS

field_block.py

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,
        )

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",
    )