Support transforming getInitialState defined via arrow functions

transforms/__testfixtures__/class/class-initial-state.input.js

@@ -311,3 +311,25 @@ var WithMultiLineType = React.createClass({
  return null;
  },
 });
+
+var WithArrowFunction = React.createClass({
+ getInitialState: (): {heyoo: number} => {
+ return {
+ heyoo: 23,
+ };
+ },
+
+ render() {
+ return null;
+ },
+});
+
+var WithArrowFunctionAndObject = React.createClass({
+ getInitialState: (): {heyoo: number} => ({
+ heyoo: 23,
+ }),
+
+ render() {
+ return null;
+ },
+});

transforms/__testfixtures__/class/class-initial-state.output.js

@@ -340,3 +340,23 @@ class WithMultiLineType extends React.Component {
  return null;
  }
 }
+
+class WithArrowFunction extends React.Component {
+ state: {heyoo: number} = {
+ heyoo: 23,
+ };
+
+ render() {
+ return null;
+ }
+}
+
+class WithArrowFunctionAndObject extends React.Component {
+ state: {heyoo: number} = {
+ heyoo: 23,
+ };
+
+ render() {
+ return null;
+ }
+}

transforms/class.js

@@ -147,13 +147,21 @@ module.exports = (file, api, options) => {
  );
 
  const hasSingleReturnStatement = value => (
- value.type === 'FunctionExpression' &&
- value.body &&
- value.body.type === 'BlockStatement' &&
- value.body.body &&
- value.body.body.length === 1 &&
- value.body.body[0].type === 'ReturnStatement' &&
- value.body.body[0].argument
+ (
+ value.type === 'ArrowFunctionExpression' &&
+ value.body &&
+ value.body.type === 'ObjectExpression'
+ ) || (
+ (
+ value.type === 'FunctionExpression' || value.type === 'ArrowFunctionExpression'
+ ) &&
+ value.body &&
+ value.body.type === 'BlockStatement' &&
+ value.body.body &&
+ value.body.body.length === 1 &&
+ value.body.body[0].type === 'ReturnStatement' &&
+ value.body.body[0].argument
+ )
  );
 
  const isInitialStateLiftable = getInitialState => {
@@ -339,7 +347,11 @@ module.exports = (file, api, options) => {
  // Collectors
  const pickReturnValueOrCreateIIFE = value => {
  if (hasSingleReturnStatement(value)) {
- return value.body.body[0].argument;
+ if (value.body.type === 'ObjectExpression') {
+ return value.body;
+ } else {
+ return value.body.body[0].argument;
+ }
  } else {
  return j.callExpression(
  value,