Allow styles to have CSS class

1 files changed, 9 insertions, 4 deletions

diff --git a/lib/html-generator.js b/lib/html-generator.js
index 8a67e53..7a03570 100644
--- a/lib/html-generator.js
+++ b/lib/html-generator.js
@@ -31,22 +31,27 @@ function HtmlGenerator() {
             pop();
         }
         for (; partIndex < style.length; partIndex++) {
-            push(style[partIndex].tagName);
+            push(style[partIndex]);
         }
     }
     
     function stylePartMatchesHtmlElement(stylePart, htmlElement) {
-        return stylePart.tagName === htmlElement;
+        return htmlElement && stylePart.tagName === htmlElement.tagName;
     }
     
     function push(element) {
         stack.push(element);
-        fragments.push(util.format("<%s>", element));
+        if (element["class"]) {
+            var attributeString = util.format(' %s="%s"', "class", element["class"]);
+        } else {
+            var attributeString = "";
+        }
+        fragments.push(util.format("<%s%s>", element.tagName, attributeString));
     }
     
     function pop() {
         var element = stack.pop();
-        fragments.push(util.format("</%s>", element));
+        fragments.push(util.format("</%s>", element.tagName));
     }
     
     function popAll() {