Skip to content

Commit

Permalink
feat: pickle support for StreamField (#10654)
Browse files Browse the repository at this point in the history
Fixes #1988
  • Loading branch information
Antoni Martyniuk authored and gasman committed Aug 21, 2023
1 parent febacf7 commit eadf9a6
Show file tree
Hide file tree
Showing 6 changed files with 65 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ Changelog
* Move `SnippetViewSet` menu registration mechanism to base `ViewSet` class (Sage Abdullah)
* Enable reference index tracking for models registered with `ModelViewSet` (Sage Abdullah)
* When copying a page or creating an alias, copy its view restrictions to the destination (Sandeep Choudhary, Suyash Singh)
* Support pickling of StreamField values (pySilver)
* Fix: Ensure that StreamField's `FieldBlock`s correctly set the `required` and `aria-describedby` attributes (Storm Heg)
* Docs: Document `WAGTAILADMIN_BASE_URL` on "Integrating Wagtail into a Django project" page (Shreshth Srivastava)
* Maintenance: Fix snippet search test to work on non-fallback database backends (Matt Westcott)
Expand Down
1 change: 1 addition & 0 deletions CONTRIBUTORS.md
Original file line number Diff line number Diff line change
Expand Up @@ -725,6 +725,7 @@
* Florent Lebreton
* Shreshth Srivastava
* Sandeep Choudhary
* Antoni Martyniuk

## Translators

Expand Down
1 change: 1 addition & 0 deletions docs/releases/5.2.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ depth: 1
* Move `SnippetViewSet` menu registration mechanism to base `ViewSet` class (Sage Abdullah)
* Enable reference index tracking for models registered with `ModelViewSet` (Sage Abdullah)
* When copying a page or creating an alias, copy its view restrictions to the destination (Sandeep Choudhary, Suyash Singh)
* Support pickling of StreamField values (pySilver)

### Bug fixes

Expand Down
26 changes: 26 additions & 0 deletions wagtail/blocks/stream_block.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@
import uuid
from collections import OrderedDict, defaultdict
from collections.abc import Mapping, MutableSequence
from pickle import PickleError

from django import forms
from django.core.exceptions import ValidationError
from django.db.models.fields import _load_field
from django.forms.utils import ErrorList
from django.utils.functional import cached_property
from django.utils.html import format_html_join
Expand Down Expand Up @@ -720,6 +722,30 @@ def __html__(self):
def __str__(self):
return self.__html__()

@staticmethod
def _deserialize_pickle_value(app_label, model_name, field_name, field_value):
"""Returns StreamValue from pickled data"""
field = _load_field(app_label, model_name, field_name)
return field.to_python(field_value)

def __reduce__(self):
try:
stream_field = self._stream_field
except AttributeError:
raise PickleError(
"StreamValue can only be pickled if it is associated with a StreamField"
)

return (
self._deserialize_pickle_value,
(
stream_field.model._meta.app_label,
stream_field.model._meta.object_name,
stream_field.name,
self.get_prep_value(),
),
)


class StreamBlockAdapter(Adapter):
js_constructor = "wagtail.blocks.StreamBlock"
Expand Down
11 changes: 11 additions & 0 deletions wagtail/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,17 @@ def deconstruct(self):
return name, path, args, kwargs

def to_python(self, value):
value = self._to_python(value)

# The top-level StreamValue is passed a reference to the StreamField, to support
# pickling. This is necessary because unpickling needs access to the StreamBlock
# definition, which cannot itself be pickled; instead we store a pointer to the
# field within the model, which gives us a path to retrieve the StreamBlock definition.

value._stream_field = self
return value

def _to_python(self, value):
if value is None or value == "":
return StreamValue(self.stream_block, [])
elif isinstance(value, StreamValue):
Expand Down
25 changes: 25 additions & 0 deletions wagtail/tests/test_streamfield.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import json
import pickle

from django.apps import apps
from django.db import connection, models
Expand All @@ -11,12 +12,14 @@
from wagtail.fields import StreamField
from wagtail.images.models import Image
from wagtail.images.tests.utils import get_test_image_file
from wagtail.models import Page
from wagtail.rich_text import RichText
from wagtail.signal_handlers import disable_reference_index_auto_update
from wagtail.test.testapp.models import (
JSONBlockCountsStreamModel,
JSONMinMaxCountStreamModel,
JSONStreamModel,
StreamPage,
)


Expand Down Expand Up @@ -600,3 +603,25 @@ def test_json_contains_lookup(self):
instance = JSONStreamModel.objects.filter(body__contains=value).first()
self.assertIsNotNone(instance)
self.assertEqual(instance.id, self.instance.id)


class TestStreamFieldPickleSupport(TestCase):
def setUp(self):
# Find root page
self.root_page = Page.objects.get(id=2)

def test_pickle_support(self):
stream_page = StreamPage(title="stream page", body=[("text", "hello")])
self.root_page.add_child(instance=stream_page)

# check that page can be serialized / deserialized
serialized = pickle.dumps(stream_page)
deserialized = pickle.loads(serialized)

# check that serialized page can be serialized / deserialized again
serialized2 = pickle.dumps(deserialized)
deserialized2 = pickle.loads(serialized2)

# check that page data is not corrupted
self.assertEqual(stream_page.body, deserialized.body)
self.assertEqual(stream_page.body, deserialized2.body)

0 comments on commit eadf9a6

Please sign in to comment.