Use HoC approach for both class- and stateless-style components

Author: sstur

src/extendReactClass.js

@@ -1,38 +1,24 @@
 import _ from 'lodash';
-import extendReactClass from './extendReactClass';
-import wrapStatelessFunction from './wrapStatelessFunction';
+import wrapComponent from './wrapComponent';
 
 /**
  * @see https://github.com/gajus/react-css-modules#options
  */
 type OptionsType = {};
 
-/**
- * Determines if the given object has the signature of a class that inherits React.Component.
- */
-const isReactComponent = (maybeReactComponent: any): boolean => {
-    return 'prototype' in maybeReactComponent && _.isFunction(maybeReactComponent.prototype.render);
-};
-
 /**
  * When used as a function.
  */
 const functionConstructor = (Component: Function, defaultStyles: Object, options: OptionsType): Function => {
-    let decoratedClass;
-
-    if (isReactComponent(Component)) {
-        decoratedClass = extendReactClass(Component, defaultStyles, options);
-    } else {
-        decoratedClass = wrapStatelessFunction(Component, defaultStyles, options);
-    }
+    let WrappedComponent = wrapComponent(Component, defaultStyles, options);
 
     if (Component.displayName) {
-        decoratedClass.displayName = Component.displayName;
+        WrappedComponent.displayName = Component.displayName;
     } else {
-        decoratedClass.displayName = Component.name;
+        WrappedComponent.displayName = Component.name;
     }
 
-    return decoratedClass;
+    return WrappedComponent;
 };
 
 /**

src/index.js

@@ -0,0 +1,70 @@
+/* eslint-disable react/prop-types */
+
+import generateAppendClassName from './generateAppendClassName';
+import hoistNonReactStatics from 'hoist-non-react-statics';
+import makeConfiguration from './makeConfiguration';
+import parseStyleName from './parseStyleName';
+import React, {Children} from 'react';
+
+/**
+ * @param {ReactElement} element
+ * @param {Object} styles
+ * @param {Object} options
+ * @returns {ReactElement}
+ */
+function updateProps(element: Object, styles: Object, options: Object) {
+    if (element == null || typeof element !== 'object') {
+        return element;
+    }
+    let {styleName, className, children} = element.props;
+    let didUpdate = false;
+    if (styleName != null) {
+        let styleNames = parseStyleName(styleName, options.allowMultiple);
+        let appendClassName = generateAppendClassName(styles, styleNames, options.errorWhenNotFound);
+        className = className ? className + ' ' + appendClassName : appendClassName;
+        didUpdate = true;
+    }
+    if (children != null) {
+        let newChildren = [];
+        let didUpdateChildren = true;
+        Children.forEach(children, (child) => {
+            let newChild = updateProps(child, styles, options);
+            newChildren.push(newChild);
+            if (newChild !== child) {
+                didUpdateChildren = true;
+            }
+        });
+        if (didUpdateChildren) {
+            children = newChildren;
+            didUpdate = true;
+        }
+    }
+    if (didUpdate) {
+        let props = {...element.props, className, children};
+        return {...element, props};
+    } else {
+        return element;
+    }
+}
+
+/**
+ * @param {Function} Component
+ * @param {Object} defaultStyles
+ * @param {Object} options
+ * @returns {ReactClass}
+ */
+export default function wrapComponent(Component: Function, defaultStyles: Object = {}, options: Object = {}) {
+    class WrappedComponent extends React.Component {
+        render() {
+            let {props} = this;
+            let {styles} = props;
+            if (styles == null) {
+                styles = defaultStyles;
+                props = {...props, styles};
+            }
+            let element = React.createElement(Component, props);
+            return updateProps(element, styles, makeConfiguration(options));
+        }
+    }
+    return hoistNonReactStatics(WrappedComponent, Component);
+}