Wrap render method created using class properties (2)

client/patch-react.js

@@ -17,8 +17,23 @@ export default (handleError = () => {}) => {
   React.createElement = function (Component, ...rest) {
     if (typeof Component === 'function') {
       const { prototype } = Component
-      if (prototype && prototype.render) {
-        prototype.render = wrapRender(prototype.render)
+
+      // assumes it's a class component if render method exists.
+      const isClassComponent = !!(prototype && prototype.render) ||
+        // subclass of React.Component or PureComponent with no render method.
+        prototype instanceof React.Component ||
+        prototype instanceof React.PureComponent
+
+      if (isClassComponent) {
+        // There's no render method in prototype
+        // when it's created with class-properties.
+        if (prototype.render) {
+          prototype.render = wrapRender(prototype.render)
+        }
+
+        // wrap the render method in runtime when the component initialized
+        // for class-properties.
+        Component = wrap(Component, withWrapOwnRender)
       } else {
         // stateless component
         Component = wrapRender(Component)
@@ -39,24 +54,42 @@ export default (handleError = () => {}) => {
   }
 
   function wrapRender (render) {
-    if (render.__wrapped) {
-      return render.__wrapped
-    }
+    return wrap(render, withHandleError)
+  }
 
-    const _render = function (...args) {
-      try {
-        return render.apply(this, args)
-      } catch (err) {
-        handleError(err)
-        return null
-      }
+  function withHandleError (fn, ...args) {
+    try {
+      return fn.apply(this, args)
+    } catch (err) {
+      handleError(err)
+      return null
     }
+  }
 
-    // copy all properties
-    Object.assign(_render, render)
+  function withWrapOwnRender (fn, ...args) {
+    const result = fn.apply(this, args)
+    if (this.render && this.hasOwnProperty('render')) {
+      this.render = wrapRender(this.render)
+    }
+    return result
+  }
+}
 
-    render.__wrapped = _render.__wrapped = _render
+function wrap (fn, around) {
+  if (fn.__wrapped) {
+    return fn.__wrapped
+  }
 
-    return _render
+  const _fn = function (...args) {
+    return around.call(this, fn, ...args)
   }
+
+  // copy all properties
+  Object.assign(_fn, fn)
+
+  _fn.prototype = fn.prototype
+
+  _fn.__wrapped = fn.__wrapped = _fn
+
+  return _fn
 }