From 759fd4c5edafd0528bfb97d784cdd4ca91a727e4 Mon Sep 17 00:00:00 2001
From: viliambalaz <balaz@sited.sk>
Date: Fri, 9 Apr 2021 12:05:06 +0200
Subject: [PATCH] #334 Create set_attributes template tag

---
 poleno/utils/templatetags/poleno/amend.py | 45 +++++++++++++
 poleno/utils/tests/test_templatetags.py   | 79 +++++++++++++++++++++++
 2 files changed, 124 insertions(+)

diff --git a/poleno/utils/templatetags/poleno/amend.py b/poleno/utils/templatetags/poleno/amend.py
index d74d3bf27..cb0a21298 100644
--- a/poleno/utils/templatetags/poleno/amend.py
+++ b/poleno/utils/templatetags/poleno/amend.py
@@ -1,5 +1,6 @@
 # vim: expandtab
 # -*- coding: utf-8 -*-
+import re
 import lxml.html
 
 from poleno.utils.template import Library
@@ -244,3 +245,47 @@ def action(fragment):
 
     context[u'_amend'].append(action)
     return u''
+
+@register.simple_tag(takes_context=True)
+def set_attributes(context, path, **kwargs):
+    u"""
+    Select elements specified by XPath and add, remove or edit their attributes.
+
+    Example:
+        {% amend %}
+          <ul>
+            <li>xxx</li>
+            <li ccc="value" ddd>yyy</li>
+            <li eee="value">zzz</li>
+          </ul>
+          {% set_attributes path=".//li[1]" aaa="value" bbb=True %}
+          {% set_attributes path=".//li[2]" ccc=None ddd=False %}
+          {% set_attributes path=".//li[3]" eee="new_value" %}
+        {% endamend %}
+
+    Result:
+        <ul>
+          <li aaa="value" bbb>xxx</li>
+          <li>yyy</li>
+          <li eee="new_value">zzz</li>
+        </ul>
+    """
+    if u'_amend' not in context:
+        context[u'_amend'] = []
+
+    def action(fragment):
+        elements = fragment.findall(path)
+        for element in elements:
+            for key, value in kwargs.items():
+                if not re.match(r'^[a-z][a-z0-9_.-]*$', key):
+                    raise ValueError(u'Invalid tag name')
+                if key not in element.attrib and not value:
+                    continue
+                elif not value:
+                    element.attrib.pop(key)
+                else:
+                    element.set(key, value if value != True else None)
+        return fragment
+
+    context[u'_amend'].append(action)
+    return u''
diff --git a/poleno/utils/tests/test_templatetags.py b/poleno/utils/tests/test_templatetags.py
index 570d295e6..017a19d04 100644
--- a/poleno/utils/tests/test_templatetags.py
+++ b/poleno/utils/tests/test_templatetags.py
@@ -423,6 +423,85 @@ def test_delete_tag(self):
                 u'</ul>'
                 u'')
 
+    def test_set_attributes_tag_add_attribute(self):
+        rendered = self._render(
+                u'{% load amend set_attributes from poleno.amend %}'
+                u'{% amend %}'
+                u'  <ul>'
+                u'    <li>xxx</li>'
+                u'    <li>yyy</li>'
+                u'    <li>zzz</li>'
+                u'  </ul>'
+                u'  {% set_attributes path=".//li" aaa="value" %}'
+                u'  {% set_attributes path=".//li[2]" bbb=True %}'
+                u'  {% set_attributes path=".//li[3]" ccc="value" ddd="value" eee=True %}'
+                u'{% endamend %}'
+                u'')
+        self.assertHTMLEqual(rendered,
+                u'<ul>'
+                u'  <li aaa="value">xxx</li>'
+                u'  <li aaa="value" bbb>yyy</li>'
+                u'  <li aaa="value" ccc="value" ddd="value" eee>zzz</li>'
+                u'</ul>'
+                u'')
+
+    def test_set_attributes_tag_edit_attribute(self):
+        rendered = self._render(
+                u'{% load amend set_attributes from poleno.amend %}'
+                u'{% amend %}'
+                u'  <ul>'
+                u'    <li aaa="value">xxx</li>'
+                u'    <li aaa="value" bbb="value" ccc="value">yyy</li>'
+                u'    <li aaa="value">zzz</li>'
+                u'  </ul>'
+                u'  {% set_attributes path=".//li" aaa="new_value" %}'
+                u'  {% set_attributes path=".//li[2]" bbb="new_value" ccc="new_value" %}'
+                u'{% endamend %}'
+                u'')
+        self.assertHTMLEqual(rendered,
+                u'<ul>'
+                u'  <li aaa="new_value">xxx</li>'
+                u'  <li aaa="new_value" bbb="new_value" ccc="new_value">yyy</li>'
+                u'  <li aaa="new_value">zzz</li>'
+                u'</ul>'
+                u'')
+
+    def test_set_attributes_tag_remove_attribute(self):
+        rendered = self._render(
+                u'{% load amend set_attributes from poleno.amend %}'
+                u'{% amend %}'
+                u'  <ul>'
+                u'    <li aaa="value">xxx</li>'
+                u'    <li aaa="value" bbb>yyy</li>'
+                u'    <li aaa="value" ccc="value" ddd="value" eee fff>zzz</li>'
+                u'  </ul>'
+                u'  {% set_attributes path=".//li" aaa=None %}'
+                u'  {% set_attributes path=".//li[2]" bbb=False %}'
+                u'  {% set_attributes path=".//li[3]" ccc=None ddd="" eee=None fff="" %}'
+                u'{% endamend %}'
+                u'')
+        self.assertHTMLEqual(rendered,
+                u'<ul>'
+                u'  <li>xxx</li>'
+                u'  <li>yyy</li>'
+                u'  <li>zzz</li>'
+                u'</ul>'
+                u'')
+
+    def test_set_attributes_tag_with_invalid_key(self):
+        with self.assertRaisesMessage(ValueError, u'Invalid tag name'):
+            rendered = self._render(
+                    u'{% load amend set_attributes from poleno.amend %}'
+                    u'{% amend %}'
+                    u'  <ul>'
+                    u'    <li>xxx</li>'
+                    u'    <li>yyy</li>'
+                    u'    <li>zzz</li>'
+                    u'  </ul>'
+                    u'  {% set_attributes path=".//li" Aaa=True %}'
+                    u'{% endamend %}'
+                    u'')
+
     def test_amend_tag_on_plain_text(self):
         rendered = self._render(
                 u'{% load amend prepend append from poleno.amend %}'