From 4c5f416e2c7d2360669e4b2b6944d8aadcfeccf8 Mon Sep 17 00:00:00 2001
From: Ethan Dalool <ethan@voussoir.net>
Date: Sun, 29 Dec 2019 13:03:44 -0800
Subject: [PATCH] Add helper function contains_class.

---
 cleanerupper/plugin.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/cleanerupper/plugin.py b/cleanerupper/plugin.py
index 6327baf..d4cb07a 100644
--- a/cleanerupper/plugin.py
+++ b/cleanerupper/plugin.py
@@ -19,15 +19,20 @@ def raise_children_and_delete(element):
         element.insert_after(children.pop(-1))
     element.decompose()
 
-def remove_class(element, cls):
-    if not hasattr(element, 'class') or element['class'] is None:
-        return
+def contains_class(element, cls):
+    try:
+        element['class']
+    except (AttributeError, KeyError):
+        return False
     if isinstance(element['class'], str):
-        if element['class'] == cls:
-            del element['class']
-            return
-        else:
-            element['class'] = element['class'].split()
+        element['class'] = element['class'].split()
+
+    return cls in element['class']
+
+def remove_class(element, cls):
+    if not contains_class(element, cls):
+        return
+
     try:
         element['class'].remove(cls)
     except IndexError: