Add support for ES6 classes in sort-comp

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');
   });
 
 });

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;

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;

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);
     }
   }

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,