From bd062ec97dcbb91a93133433c8e39c9c0ecdebf9 Mon Sep 17 00:00:00 2001
From: Sylvain Boissel <sylvain@boissel.dev>
Date: Thu, 13 Jun 2024 12:10:13 +0200
Subject: [PATCH] Improve the new filter

---
 .../content_manager/blocks/callout.html       |  2 +-
 .../content_manager/blocks/highlight.html     |  6 +----
 .../templatetags/wagtail_dsfr_tags.py         | 25 +++++++++++--------
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/content_manager/templates/content_manager/blocks/callout.html b/content_manager/templates/content_manager/blocks/callout.html
index 0ecfd0cb..12de18e8 100644
--- a/content_manager/templates/content_manager/blocks/callout.html
+++ b/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 %}
diff --git a/content_manager/templates/content_manager/blocks/highlight.html b/content_manager/templates/content_manager/blocks/highlight.html
index 0e0e6a16..6149fbfe 100644
--- a/content_manager/templates/content_manager/blocks/highlight.html
+++ b/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>
diff --git a/content_manager/templatetags/wagtail_dsfr_tags.py b/content_manager/templatetags/wagtail_dsfr_tags.py
index 34eafa7d..b5d994df 100644
--- a/content_manager/templatetags/wagtail_dsfr_tags.py
+++ b/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))