Allow register_snippet() to be passed a SnippetViewSet as the only argument

wagtail/snippets/models.py

@@ -69,33 +69,48 @@ class SnippetAdminURLFinder:
             )
 
 
-def register_snippet(model, viewset=None):
+def register_snippet(registerable, viewset=None):
     if DEFER_REGISTRATION:
         # Models may not have been fully loaded yet, so defer registration until they are -
         # add it to the list of registrations to be processed by register_deferred_snippets
-        DEFERRED_REGISTRATIONS.append((model, viewset))
+        DEFERRED_REGISTRATIONS.append((registerable, viewset))
     else:
-        _register_snippet_immediately(model, viewset)
+        _register_snippet_immediately(registerable, viewset)
 
-    return model
+    return registerable
 
 
-def _register_snippet_immediately(model, viewset=None):
+def _register_snippet_immediately(registerable, viewset=None):
     # Register the viewset and formfield for this snippet model,
     # skipping the check for whether models are loaded
 
+    from wagtail.snippets.views.snippets import SnippetViewSet
+
+    if isinstance(registerable, str):
+        registerable = import_string(registerable)
+    if isinstance(viewset, str):
+        viewset = import_string(viewset)
+
+    if issubclass(registerable, SnippetViewSet):
+        # register_snippet(CustomViewSet) or
+        # @register_snippet on class CustomViewSet
+        # Note: the viewset class must define a `model` attribute
+        viewset = registerable
+        admin_viewset = viewset()
+        model = admin_viewset.model
+    else:
+        # register_snippet(SnippetModel, viewset=CustomViewSet) or
+        # register_snippet(SnippetModel) or
+        # @register_snippet on class SnippetModel
+        if viewset is None:
+            viewset = SnippetViewSet
+        model = registerable
+        admin_viewset = viewset(model)
+
     if model in SNIPPET_MODELS:
         # Do not create duplicate registrations of the same model
         return
 
-    from wagtail.snippets.views.snippets import SnippetViewSet
-
-    if viewset is None:
-        viewset = SnippetViewSet
-    elif isinstance(viewset, str):
-        viewset = import_string(viewset)
-
-    admin_viewset = viewset(model)
     viewsets.register(admin_viewset)
 
     SNIPPET_MODELS.append(model)
@@ -109,8 +124,8 @@ def register_deferred_snippets():
     """
     global DEFER_REGISTRATION
     DEFER_REGISTRATION = False
-    for model, viewset in DEFERRED_REGISTRATIONS:
-        _register_snippet_immediately(model, viewset)
+    for registerable, viewset in DEFERRED_REGISTRATIONS:
+        _register_snippet_immediately(registerable, viewset)
 
 
 def create_extra_permissions(*args, using=DEFAULT_DB_ALIAS, **kwargs):

wagtail/snippets/tests/test_viewset.py

@@ -2,6 +2,7 @@ from django.contrib.admin.utils import quote
 from django.contrib.auth import get_permission_codename
 from django.contrib.auth.models import Permission
 from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import ImproperlyConfigured
 from django.test import TestCase, TransactionTestCase
 from django.urls import reverse
 from django.utils.timezone import now
@@ -14,6 +15,8 @@ from wagtail.blocks.field_block import FieldBlockAdapter
 from wagtail.coreutils import get_dummy_request
 from wagtail.models import Locale, Workflow, WorkflowContentType
 from wagtail.snippets.blocks import SnippetChooserBlock
+from wagtail.snippets.models import register_snippet
+from wagtail.snippets.views.snippets import SnippetViewSet
 from wagtail.snippets.widgets import AdminSnippetChooser
 from wagtail.test.testapp.models import (
     Advert,
@@ -25,6 +28,17 @@ from wagtail.test.testapp.models import (
 from wagtail.test.utils import WagtailTestUtils
 
 
+class TestIncorrectRegistration(TestCase):
+    def test_no_model_set_or_passed(self):
+        # The base SnippetViewSet class has no `model` attribute set,
+        # so using it directly should raise an error
+        with self.assertRaisesMessage(
+            ImproperlyConfigured,
+            "SnippetViewSet must be passed a model or define a model attribute.",
+        ):
+            register_snippet(SnippetViewSet)
+
+
 class BaseSnippetViewSetTests(WagtailTestUtils, TestCase):
     model = None
 

wagtail/test/testapp/wagtail_hooks.py

@@ -301,6 +301,8 @@ class DraftStateModelViewSet(SnippetViewSet):
 
 
 class ModeratedModelViewSet(SnippetViewSet):
+    model = ModeratedModel
+
     list_filter = {
         "text": ["exact", "contains"],
         "first_published_at": ["exact", "lt", "gt"],
@@ -310,4 +312,4 @@ class ModeratedModelViewSet(SnippetViewSet):
 register_snippet(FullFeaturedSnippet, viewset=FullFeaturedSnippetViewSet)
 register_snippet(RevisableChildModel, viewset=RevisableChildModelViewSet)
 register_snippet(DraftStateModel, viewset=DraftStateModelViewSet)
-register_snippet(ModeratedModel, viewset=ModeratedModelViewSet)
+register_snippet(ModeratedModelViewSet)