website/website/common/streamfield.py

96 lines
2.3 KiB
Python
Raw Normal View History

2022-06-26 19:25:30 +01:00
from typing import Iterator
2022-06-26 18:37:04 +01:00
from django.utils import lorem_ipsum
2022-06-26 19:25:30 +01:00
from django.utils.html import format_html_join, strip_tags
from django.utils.text import smart_split
2022-06-26 18:37:04 +01:00
from wagtail import blocks
from wagtail.embeds.blocks import EmbedBlock
2022-06-27 20:40:55 +01:00
from wagtail.images.blocks import ImageChooserBlock
2022-06-26 18:37:04 +01:00
2022-06-26 19:52:20 +01:00
IGNORE_PLAINTEXT_BLOCKS = (blocks.RawHTMLBlock, EmbedBlock)
2022-06-26 19:25:30 +01:00
2022-06-27 19:58:08 +01:00
RICH_TEXT_FEATURES = [
"h1",
"h2",
"h3",
"h4",
"h5",
"h6",
"bold",
"italic",
"ol",
"ul",
"link",
"document-link",
"code",
"strikethrough",
]
2022-06-27 20:40:55 +01:00
RICH_TEXT_FEATURES_SIMPLE = [
"bold",
"italic",
"link",
"document-link",
"code",
"strikethrough",
]
2022-06-26 18:37:04 +01:00
class LoremBlock(blocks.StructBlock):
paragraphs = blocks.IntegerBlock(min_value=1)
2022-06-26 19:25:30 +01:00
def render(self, value: dict, context: dict | None = None) -> str:
2022-06-26 18:37:04 +01:00
return format_html_join(
"\n\n",
"<p>{}</p>",
[(paragraph,) for paragraph in lorem_ipsum.paragraphs(value["paragraphs"])],
)
class Meta:
icon = "openquote"
label = "Lorem Ipsum"
2022-06-27 20:40:55 +01:00
class ImageCaptionBlock(blocks.StructBlock):
image = ImageChooserBlock()
caption = blocks.RichTextBlock(features=RICH_TEXT_FEATURES_SIMPLE)
class Meta:
icon = "image"
label = "Image with caption"
template = "common/blocks/image-caption.html"
2022-06-26 18:37:04 +01:00
def get_blocks() -> list[tuple[str, blocks.BaseBlock]]:
return [
("embed", EmbedBlock()),
2022-06-27 19:58:08 +01:00
("rich_text", blocks.RichTextBlock(features=RICH_TEXT_FEATURES)),
2022-06-26 18:37:04 +01:00
("lorem", LoremBlock()),
("html", blocks.RawHTMLBlock()),
2022-06-27 20:40:55 +01:00
("image", ImageCaptionBlock()),
2022-06-26 18:37:04 +01:00
]
2022-06-26 19:25:30 +01:00
def get_plain_text(value: blocks.StreamValue) -> Iterator[str]:
for block in value:
2022-06-26 19:52:20 +01:00
if isinstance(block.block_type, IGNORE_PLAINTEXT_BLOCKS):
2022-06-26 19:25:30 +01:00
continue
yield strip_tags(str(block))
2022-06-26 19:52:20 +01:00
def truncate_streamfield(value: blocks.StreamValue, words: int) -> str:
collected_words: list[str] = []
for block_text in get_plain_text(value):
collected_words.extend(smart_split(block_text))
if len(collected_words) >= words:
break
return " ".join(collected_words[:words])
2022-06-26 19:25:30 +01:00
def get_word_count(value: blocks.StreamValue) -> int:
count = 0
for chunk in get_plain_text(value):
count += len(list(smart_split(chunk)))
return count