mbasso
diff --git a/‎README.md
+2-2 b/‎README.md
+2-2
diff --git a/‎src/classInjector.js
-35 b/‎src/classInjector.js
-35
diff --git a/‎src/index.js
+6-3 b/‎src/index.js
+6-3
diff --git a/‎src/injector.js
+75 b/‎src/injector.js
+75
diff --git a/‎test/index.js renamed to ‎test/DOMUtils.js
+1-1 b/‎test/index.js renamed to ‎test/DOMUtils.js
+1-1
diff --git a/‎test/classInjector.js
-118 b/‎test/classInjector.js
-118
@@ -31,7 +31,7 @@ Using ReactCSSOM means:
 - No limitations to CSS power
 - Easy to use with CSS Preprocessors (no more stuff needed)
 - Highly expressive
-- Lightweight, only 50 lines of js
+- Lightweight, only 100 lines of js
 
 ## Installation
 
@@ -171,7 +171,7 @@ and here is the corresponding css, note the global selector:
 
 ### Additional API
 
-ReactCSSOM also exposes 2 simple API to mount and unmount styles, here they are:
+ReactCSSOM also exposes 2 simple API to mount and unmount styles in the DOM, here they are:
 
 ```js
 import ReactCSSOM from 'react-cssom';
 
@@ -1,8 +1,11 @@
-import initClassInjector from './classInjector';
-
-initClassInjector();
+import React from 'react';
+import inject from './injector';
 
 const ReactCSSOM = {
+  initInjector: () => {
+    React.Component = inject(React.Component)();
+  },
+  inject,
   mount: (css) => {
     const head = document.head || document.getElementsByTagName('head')[0];
     const style = document.createElement('style');
 
@@ -0,0 +1,75 @@
+import React from 'react';
+
+const getReactCSSOMClass = (componentName) => `⚛${componentName}`;
+
+function injectInElement(component, componentName) {
+  let result = component;
+  if (component) {
+    const classToInject = getReactCSSOMClass(componentName);
+    if (typeof component.type === 'string') {
+      result = React.cloneElement(
+        component,
+        {
+          className: `${classToInject} ${component.props.className || ''}`,
+        }
+      );
+    } else if (typeof component.type === 'function') {
+      result = (
+        <div className={`${classToInject}`}>
+          {component}
+        </div>
+			);
+    }
+  }
+  return result;
+}
+
+function injectInComponent(Component, componentName) {
+  const componentPrototype = Component.prototype;
+  const componentConstructor = Component;
+
+  //eslint-disable-next-line
+  const newComponent = function ReactComponent(...params) {
+    componentConstructor.apply(this, [...params]);
+    const originalRender = this.render;
+    this.render = () => (
+      injectInElement(
+        originalRender.apply(this),
+        componentName || this.constructor.displayName || this.constructor.name
+      )
+    );
+  };
+
+  //eslint-disable-next-line
+  newComponent.prototype = componentPrototype;
+  Object.keys(Component).forEach((x) => {
+    newComponent[x] = Component[x];
+  });
+
+  return newComponent;
+}
+
+function injectInFunctionalComponent(component, componentName) {
+  return (props) => injectInElement(
+    component(props),
+    componentName
+  );
+}
+
+export default function inject(component) {
+  return (componentName) => {
+    let result = null;
+    if (
+      typeof component === 'function' &&
+      component.prototype &&
+      component.prototype.isReactComponent
+    ) {
+      result = injectInComponent(component, componentName);
+    } else if (typeof component === 'function') {
+      result = injectInFunctionalComponent(component, componentName);
+    } else if (component !== null && typeof component === 'object') {
+      result = injectInElement(component, componentName);
+    }
+    return result;
+  };
+}
@@ -2,7 +2,7 @@ import expect from 'expect';
 import { jsdom } from 'jsdom';
 import ReactCSSOM from '../src/index';
 
-describe('Class', () => {
+describe('DOMUtils', () => {
   afterEach(() => {
     global.document = jsdom('<!doctype html><html><head></head><body></body></html>');
   });