Improve the new filter

content_manager/templates/content_manager/blocks/callout.html

@@ -5,7 +5,7 @@
     {{ value.title }}
     </{{ value.heading_tag | default:"p" }}>
   {% endif %}
-  {{ value.text | richtext | richtext_p_add_class:"fr-callout__text" }}
+  {{ value.text | richtext_p_add_class:"fr-callout__text" }}
   {% if value.button.url %}
     {% include "content_manager/blocks/button.html" with button=value.button %}
   {% endif %}

content_manager/templates/content_manager/blocks/highlight.html

@@ -1,9 +1,5 @@
 {% load wagtailcore_tags wagtail_dsfr_tags %}
 
 <div class="fr-highlight{% if value.color %} fr-highlight--{{ value.color }}{% endif %}">
-  {% if value.size %}
-    {{ value.text | richtext | richtext_p_add_class:value.size }}
-  {% else %}
-    {{ value.text | richtext }}
-  {% endif %}
+  {{ value.text | richtext_p_add_class:value.size }}
 </div>

content_manager/templatetags/wagtail_dsfr_tags.py

@@ -1,7 +1,8 @@
+from bs4 import BeautifulSoup
 from django import template
 from django.conf import settings
 from django.template.context import Context
-from django.utils.html import SafeString, mark_safe
+from django.utils.html import mark_safe
 from wagtail.rich_text import RichText
 
 from content_manager.models import MegaMenu
@@ -30,18 +31,22 @@ def richtext_p_add_class(value, class_name: str):
     """
     Adds a CSS class to a Richtext-generated paragraph.
 
-    Intended to be used right after a `| richext` filter
+    Intended to be used right after a `| richext` filter in case of a RichTextField
+    (not necessary for a RichTextBlock)
     """
 
     if not class_name:
-        raise ValueError("Missing or empty parameter: class_name.")
+        return value
 
     if isinstance(value, RichText):
-        # In case of a RichTextBlock, render it
-        value = value.__html__()
+        # In case of a RichTextBlock, first render it
+        value = str(value)
 
-    if isinstance(value, SafeString):
-        # In case of a RichTextField, of after rendering a RichTextBlock
-        return mark_safe(value.replace("<p data-block-key", f'<p class="{class_name}" data-block-key'))
-    else:
-        return value
+    soup = BeautifulSoup(value, "html.parser")
+
+    paragraphs = soup.find_all("p")
+
+    for p in paragraphs:
+        p["class"] = p.get("class", []) + [class_name]
+
+    return mark_safe(str(soup))