From 898ea7613fbee0269f63185f187e2889ca1da843 Mon Sep 17 00:00:00 2001
From: Joe Lencioni <joe.lencioni@airbnb.com>
Date: Wed, 3 Feb 2016 17:46:46 -0800
Subject: [PATCH] Add support for ES6 classes in sort-comp

In a comment in sort-comp, I noticed the following comment:

> NOTE: only works on React.createClass() syntax, not ES6 class.

I thought that it would be really lovely if this worked with ES6
classes, so I decided to take a crack at adding support.

Fixes #27
---
 test/__tests__/sort-comp-test.js |  2 ++
 test/sort-comp-test2.js          | 35 ++++++++++++++++++++++++++++++++
 test/sort-comp-test2.output.js   | 35 ++++++++++++++++++++++++++++++++
 transforms/sort-comp.js          | 27 +++++++++++++++++++-----
 transforms/utils/ReactUtils.js   |  3 +++
 5 files changed, 97 insertions(+), 5 deletions(-)
 create mode 100644 test/sort-comp-test2.js
 create mode 100644 test/sort-comp-test2.output.js

diff --git a/test/__tests__/sort-comp-test.js b/test/__tests__/sort-comp-test.js
index d26d780f..b661d268 100644
--- a/test/__tests__/sort-comp-test.js
+++ b/test/__tests__/sort-comp-test.js
@@ -4,6 +4,8 @@ describe('sort-comp', () => {
 
   it('transforms correctly', () => {
     test('sort-comp', 'sort-comp-test');
+
+    test('sort-comp', 'sort-comp-test2');
   });
 
 });
diff --git a/test/sort-comp-test2.js b/test/sort-comp-test2.js
new file mode 100644
index 00000000..959267cd
--- /dev/null
+++ b/test/sort-comp-test2.js
@@ -0,0 +1,35 @@
+var React = require('react/addons');
+
+const propTypes = {};
+
+// comment above class
+class MyComponent extends React.Component {
+  // comment at top of createClass
+  // this will be attached to first method
+
+  render() {
+    return <div />;
+  }
+
+  // comment on componentDidMount
+  componentDidMount() {
+  }
+
+  renderFoo() {
+    // other render* function
+  }
+
+  renderBar() {
+    // should come before renderFoo
+  }
+
+  myOwnMethod(foo) {
+    // comment within method
+  }
+
+}
+
+MyComponent.propTypes = propTypes;
+
+/* comment at end */
+module.exports = MyComponent;
diff --git a/test/sort-comp-test2.output.js b/test/sort-comp-test2.output.js
new file mode 100644
index 00000000..af087bde
--- /dev/null
+++ b/test/sort-comp-test2.output.js
@@ -0,0 +1,35 @@
+var React = require('react/addons');
+
+const propTypes = {};
+
+// comment above class
+class MyComponent extends React.Component {
+  // comment on componentDidMount
+  componentDidMount() {
+  }
+
+  myOwnMethod(foo) {
+    // comment within method
+  }
+
+  renderBar() {
+    // should come before renderFoo
+  }
+
+  renderFoo() {
+    // other render* function
+  }
+
+  // comment at top of createClass
+  // this will be attached to first method
+
+  render() {
+    return <div />;
+  }
+
+}
+
+MyComponent.propTypes = propTypes;
+
+/* comment at end */
+module.exports = MyComponent;
diff --git a/transforms/sort-comp.js b/transforms/sort-comp.js
index 82a50fc0..5b0490ce 100644
--- a/transforms/sort-comp.js
+++ b/transforms/sort-comp.js
@@ -13,8 +13,6 @@
  *      'render'
  *    ]
  *  }],
- *
- * NOTE: only works on React.createClass() syntax, not ES6 class.
  */
 module.exports = function(fileInfo, api, options) {
   const j = api.jscodeshift;
@@ -52,14 +50,33 @@ module.exports = function(fileInfo, api, options) {
     }
   };
 
+  const sortClassProperties = classPath => {
+    const spec = ReactUtils.getClassExtendReactSpec(classPath);
+
+    if (spec) {
+      spec.body.sort(propertyComparator);
+    }
+  };
+
   if (
     options['explicit-require'] === false ||
     ReactUtils.hasReact(root)
   ) {
-    const sortCandidates = ReactUtils.findReactCreateClass(root);
+    const createClassSortCandidates = ReactUtils.findReactCreateClass(root);
+    const es6ClassSortCandidates = ReactUtils.findReactES6ClassDeclaration(root);
+
+    if (createClassSortCandidates.size() > 0) {
+      createClassSortCandidates.forEach(sortComponentProperties);
+    }
+
+    if (es6ClassSortCandidates.size() > 0) {
+      es6ClassSortCandidates.forEach(sortClassProperties);
+    }
 
-    if (sortCandidates.size() > 0) {
-      sortCandidates.forEach(sortComponentProperties);
+    if (
+      createClassSortCandidates.size() > 0 ||
+      es6ClassSortCandidates.size() > 0
+    ) {
       return root.toSource(printOptions);
     }
   }
diff --git a/transforms/utils/ReactUtils.js b/transforms/utils/ReactUtils.js
index 7cdaf7fc..970d2a77 100644
--- a/transforms/utils/ReactUtils.js
+++ b/transforms/utils/ReactUtils.js
@@ -133,6 +133,8 @@ module.exports = function(j) {
     }
   };
 
+  const getClassExtendReactSpec = classPath => classPath.value.body;
+
   const createCreateReactClassCallExpression = properties =>
     j.callExpression(
       j.memberExpression(
@@ -155,6 +157,7 @@ module.exports = function(j) {
     findReactCreateClassExportDefault,
     getComponentName,
     getReactCreateClassSpec,
+    getClassExtendReactSpec,
     hasMixins,
     hasModule,
     hasReact,