Allow styles to have CSS class

2 files changed, 11 insertions, 5 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() {
diff --git a/lib/styles.js b/lib/styles.js
index 06d63a8..8120a76 100644
--- a/lib/styles.js
+++ b/lib/styles.js
@@ -22,6 +22,7 @@ function element(tagName, options) {
     options = options || {};
     return {
         tagName: tagName,
-        fresh: options.fresh
+        fresh: options.fresh,
+        "class": options["class"]
     };
 }