From ae9cf8435f926cab0131298a9ba8a680167268c7 Mon Sep 17 00:00:00 2001
From: Cale Newman <newman.cale@gmail.com>
Date: Wed, 13 Jul 2016 03:43:03 -0500
Subject: [PATCH] extend the existing babel plugin to transform class property
 arrow functions into hot-reloadable class methods (#242)

---
 src/babel/index.js                            | 34 +++++++++++++++++
 test/AppContainer/.babelrc                    |  4 ++
 test/babel/fixtures/class-properties/.babelrc |  4 ++
 .../babel/fixtures/class-properties/actual.js |  5 +++
 .../fixtures/class-properties/expected.js     | 38 +++++++++++++++++++
 5 files changed, 85 insertions(+)
 create mode 100644 test/AppContainer/.babelrc
 create mode 100644 test/babel/fixtures/class-properties/.babelrc
 create mode 100644 test/babel/fixtures/class-properties/actual.js
 create mode 100644 test/babel/fixtures/class-properties/expected.js

diff --git a/src/babel/index.js b/src/babel/index.js
index fb9bf1efe..55a319050 100644
--- a/src/babel/index.js
+++ b/src/babel/index.js
@@ -14,6 +14,16 @@ const buildTagger = template(`
 })();
 `);
 
+const buildNewClassProperty = (t, key, identifier, params) => {
+  const returnExpression = t.callExpression(
+    t.memberExpression(t.thisExpression(), identifier),
+    params
+  );
+  const blockStatement = t.blockStatement([t.returnStatement(returnExpression)]);
+  const newArrowFunction = t.arrowFunctionExpression(params, blockStatement);
+  return t.classProperty(key, newArrowFunction);
+};
+
 module.exports = function plugin(args) {
   // This is a Babel plugin, but the user put it in the Webpack config.
   if (this && this.callback) {
@@ -125,6 +135,30 @@ module.exports = function plugin(args) {
           node.body.push(buildSemi());
         },
       },
+
+      Class(classPath) {
+        const classBody = classPath.get('body');
+
+        classBody.get('body').forEach(path => {
+          if (path.isClassProperty()) {
+            const { node } = path;
+
+            if (node.value.type === 'ArrowFunctionExpression') {
+              const { params } = node.value;
+              const newIdentifier = t.identifier(`__${node.key.name}__REACT_HOT_LOADER__`);
+
+              // create a new method on the class that the original class property function
+              // calls, since the method is able to be replaced by RHL
+              const newMethod = t.classMethod('method', newIdentifier, params, node.value.body);
+              path.insertAfter(newMethod);
+
+              // replace the original class property function with a function that calls
+              // the new class method created above
+              path.replaceWith(buildNewClassProperty(t, node.key, newIdentifier, params));
+            }
+          }
+        });
+      },
     },
   };
 };
diff --git a/test/AppContainer/.babelrc b/test/AppContainer/.babelrc
new file mode 100644
index 000000000..f91228920
--- /dev/null
+++ b/test/AppContainer/.babelrc
@@ -0,0 +1,4 @@
+{
+  "presets": ["es2015", "stage-1", "react"],
+  "plugins": ["../../src/babel"]
+}
diff --git a/test/babel/fixtures/class-properties/.babelrc b/test/babel/fixtures/class-properties/.babelrc
new file mode 100644
index 000000000..665cc9b04
--- /dev/null
+++ b/test/babel/fixtures/class-properties/.babelrc
@@ -0,0 +1,4 @@
+{
+  "presets": ["es2015", "stage-1"],
+  "plugins": ["../../../../src/babel"]
+}
diff --git a/test/babel/fixtures/class-properties/actual.js b/test/babel/fixtures/class-properties/actual.js
new file mode 100644
index 000000000..c0c0df9c0
--- /dev/null
+++ b/test/babel/fixtures/class-properties/actual.js
@@ -0,0 +1,5 @@
+class Stuff {
+  foo = (a, b) => {
+    return a(b);
+  };
+}
diff --git a/test/babel/fixtures/class-properties/expected.js b/test/babel/fixtures/class-properties/expected.js
new file mode 100644
index 000000000..6af90a68b
--- /dev/null
+++ b/test/babel/fixtures/class-properties/expected.js
@@ -0,0 +1,38 @@
+"use strict";
+
+var _createClass = function () { function defineProperties(target, props) { for (var i = 0; i < props.length; i++) { var descriptor = props[i]; descriptor.enumerable = descriptor.enumerable || false; descriptor.configurable = true; if ("value" in descriptor) descriptor.writable = true; Object.defineProperty(target, descriptor.key, descriptor); } } return function (Constructor, protoProps, staticProps) { if (protoProps) defineProperties(Constructor.prototype, protoProps); if (staticProps) defineProperties(Constructor, staticProps); return Constructor; }; }();
+
+function _classCallCheck(instance, Constructor) { if (!(instance instanceof Constructor)) { throw new TypeError("Cannot call a class as a function"); } }
+
+var Stuff = function () {
+  function Stuff() {
+    var _this = this;
+
+    _classCallCheck(this, Stuff);
+
+    this.foo = function (a, b) {
+      return _this.__foo__REACT_HOT_LOADER__(a, b);
+    };
+  }
+
+  _createClass(Stuff, [{
+    key: "__foo__REACT_HOT_LOADER__",
+    value: function __foo__REACT_HOT_LOADER__(a, b) {
+      return a(b);
+    }
+  }]);
+
+  return Stuff;
+}();
+
+;
+
+(function () {
+  if (typeof __REACT_HOT_LOADER__ === 'undefined') {
+    return;
+  }
+
+  __REACT_HOT_LOADER__.register(Stuff, "Stuff", "/Users/calenewman/code/react-hot-loader/test/babel/fixtures/class-properties/actual.js");
+})();
+
+;