fix(Assertions): use React

src/assertions.js

@@ -1,36 +1,19 @@
 import * as t from 'babel-types'
 
-const expectStatic = (key, isStatic) => !isStatic && key.buildCodeFrameError(`'${key} must be defined as static`)
-
-export const isHandledAssignment = (left, right, property) => {
+export const isHandledAssignment = (path, { left, right, property }) => {
   if (!t.isMemberExpression(left) || !t.isIdentifier(property, { name: 'handledProps' })) return false
-  if (!t.isArrayExpression(right)) right.buildCodeFrameError('`handledProps` must be an array')
+  if (!t.isArrayExpression(right)) throw path.buildCodeFrameError('`handledProps` must be an array')
 
   return true
 }
 
-export const isHandledProperty = (key, value, isStatic) => {
-  if (!t.isIdentifier(key, { name: 'handledProps' })) return false
-  if (!t.isArrayExpression(value)) value.buildCodeFrameError('`handledProps` must be an array')
-  expectStatic(key, isStatic)
-
-  return true
-}
 
-export const isPropsAssignment = (left, right, property) => {
+export const isPropsAssignment = (path, { left, right, property }) => {
   if (!t.isMemberExpression(left)) return false
   if (!t.isIdentifier(property, { name: 'defaultProps' }) && !t.isIdentifier(property, { name: 'propTypes' })) {
     return false
   }
-  if (!t.isObjectExpression(right)) right.buildCodeFrameError('`defaultProps` and `propTypes` must be an array')
-
-  return true
-}
-
-export const isPropsProperty = (key, value, isStatic) => {
-  if (!t.isIdentifier(key, { name: 'defaultProps' }) && !t.isIdentifier(key, { name: 'propTypes' })) return false
-  if (!t.isObjectExpression(value)) value.buildCodeFrameError('`defaultProps` and `propTypes` must be object')
-  expectStatic(key, isStatic)
+  if (!t.isObjectExpression(right)) throw path.buildCodeFrameError('`defaultProps` and `propTypes` must be an object')
 
   return true
 }

src/index.js

@@ -1,68 +1,36 @@
 import {
-  isHandledAssignment,
-  isHandledProperty,
-  isPropsAssignment,
-  isPropsProperty,
-} from './assertions'
-import { findClassIdentifier, generateExpression } from './helpers'
-import Store from './Store'
+  entryVisitor,
+  importVisitor,
+  propVisitor,
+} from './visitors'
+import {
+  createPropertyExpression,
+  insertAfterPath,
+  Store,
+} from './util'
+
+const insertEntries = entries => entries.forEach(({ identifier, path, props }) => {
+  insertAfterPath(path, createPropertyExpression(identifier, props))
+})
 
-export default function ({ types: t }) {
+const plugin = () => {
   return {
+    pre() {
+      this.store = new Store()
+    },
     visitor: {
       Program(programPath) {
-        const store = new Store()
-
-        programPath.traverse({
-          AssignmentExpression(path) {
-            const { left, right } = path.node
-            const { object, property } = left
-            const { name: identifier } = object
-
-            if (isHandledAssignment(left, right, property)) {
-              const { elements } = right
-
-              elements.forEach(element => store.add(identifier, element.value))
-              path.remove()
-
-              return
-            }
+        programPath.traverse(importVisitor, this.store)
 
-            if (isPropsAssignment(left, right, property)) {
-              const { properties } = right
-              properties.forEach(item => store.add(identifier, item.key.name))
-            }
-          },
-          ClassProperty(path) {
-            const { key, value } = path.node
+        if (!this.store.hasImport) return
 
-            if (isHandledProperty(key, value, path.node.static)) {
-              const { elements } = value
+        programPath.traverse(entryVisitor, this.store)
+        programPath.traverse(propVisitor, this.store)
 
-              elements.forEach(element => store.add(findClassIdentifier(path), element.value))
-              path.remove()
-
-              return
-            }
-
-            if (isPropsProperty(key, value, path.node.static)) {
-              const { properties } = value
-              properties.forEach(property => store.add(findClassIdentifier(path), property.key.name))
-            }
-          },
-        })
-
-        programPath.traverse({
-          'ClassDeclaration|FunctionDeclaration'(path) {
-            const { name } = path.node.id
-
-            if (!store.has(name)) return
-            if (t.isExportDeclaration(path.parentPath)) path = path.parentPath
-
-            path.insertAfter(generateExpression(store, name))
-          },
-        })
+        insertEntries(this.store.getEntries())
       },
     },
   }
 }
+
+export default plugin

src/util/Store.js

@@ -0,0 +1,20 @@
+import _ from 'lodash'
+
+export default class Store {
+  entries = {}
+  hasImport = false
+
+  addProps = (name, newProps) => {
+    this.entries[name].props = _.union(this.entries[name].props, newProps)
+  }
+
+  createEntry = (name, path) => {
+    this.entries[name] = { path, props: [] }
+  }
+
+  getEntries = () => _.map(this.entries, ({ path, props }, identifier) => {
+    return { identifier, path, props }
+  })
+
+  hasEntry = name => _.has(this.entries, name)
+}

src/util/createPropertyExpression.js

@@ -0,0 +1,13 @@
+import _ from 'lodash'
+import * as t from 'babel-types'
+
+const createPropertyExpression = (identifier, props) => {
+  const entries = _.uniq(props).sort().map(prop => t.stringLiteral(prop))
+
+  const left = t.memberExpression(t.identifier(identifier), t.identifier('handledProps'))
+  const right = t.arrayExpression(entries)
+
+  return t.expressionStatement(t.assignmentExpression('=', left, right))
+}
+
+export default createPropertyExpression

src/util/getBody.js

@@ -0,0 +1,3 @@
+const getBody = ({ node: { body } }) => body.body
+
+export default getBody

src/util/getClassDeclaration.js

@@ -0,0 +1,5 @@
+import * as t from 'babel-types'
+
+const getClassDeclaration = path => path.findParent(parentPath => t.isClassDeclaration(parentPath))
+
+export default getClassDeclaration

src/util/getEntryIdentifier.js

@@ -0,0 +1,21 @@
+import * as t from 'babel-types'
+
+const findVariableDeclarator = path => path.findParent(parentPath => t.isVariableDeclarator(parentPath))
+
+const getName = ({ node: { id: { name } } }) => name
+
+const getFunctionIdentifier = path => {
+  if (t.isFunctionDeclaration(path)) return getName(path)
+  if (t.isArrowFunctionExpression(path) || t.isFunctionExpression(path)) return getName(findVariableDeclarator(path))
+
+  throw path.buildCodeFrameError('`path` is unsupported Function definition')
+}
+
+const getEntryIdentifier = path => {
+  if (t.isClass(path)) return getName(path)
+  if (t.isFunction(path)) return getFunctionIdentifier(path)
+
+  throw path.buildCodeFrameError('`path` must be Class or Function definition')
+}
+
+export default getEntryIdentifier

src/util/getExpressionIdentifier.js

@@ -0,0 +1,3 @@
+const getExpressionIdentifier = ({ node: { left: { object } } }) => object.name
+
+export default getExpressionIdentifier

src/util/hasReturnStatement.js

@@ -0,0 +1,8 @@
+import * as t from 'babel-types'
+
+const hasReturnStatement = body => {
+  if (t.isBlockStatement(body)) return body.some(member => t.isReturnStatement(member))
+  return !!body
+}
+
+export default hasReturnStatement

src/util/index.js

@@ -0,0 +1,20 @@
+export { default as createPropertyExpression } from './createPropertyExpression'
+export { default as insertAfterPath } from './insertAfterPath'
+
+export { default as getClassDeclaration } from './getClassDeclaration'
+export { default as getExpressionIdentifier } from './getExpressionIdentifier'
+export { default as getEntryIdentifier } from './getEntryIdentifier'
+
+export { default as isReactClass } from './isReactClass'
+export { default as isReactFunction } from './isReactFunction'
+export { default as isReactImport } from './isReactImport'
+
+export { default as isValidExpression } from './isValidExpression'
+export { default as isValidProperty } from './isValidProperty'
+
+export { default as Store } from './Store'
+export {
+  isArrayValue,
+  isObjectValue,
+  isStaticProperty,
+} from './types'

src/util/insertAfterPath.js

@@ -0,0 +1,14 @@
+import * as t from 'babel-types'
+
+const insertAfterPath = (path, expression) => {
+  if (t.isExportDeclaration(path.parent)) {
+    const parent = path.findParent(parentPath => parentPath.isExportDeclaration())
+
+    parent.insertAfter(expression)
+    return
+  }
+
+  path.insertAfter(expression)
+}
+
+export default insertAfterPath

src/util/isReactClass.js

@@ -0,0 +1,20 @@
+import * as t from 'babel-types'
+
+import getBody from './getBody'
+import hasReturnStatement from './hasReturnStatement'
+
+const getRenderMethod = path => getBody(path).find(member => {
+  return t.isClassMethod(member) && t.isIdentifier(member.key, { name: 'render' })
+})
+
+const hasSuperClass = ({ node: { superClass } }) => !!superClass
+
+const hasValidRenderMethod = renderMethod => !!renderMethod && hasReturnStatement(renderMethod)
+
+const isClass = path => t.isClassDeclaration(path) || t.isClassExpression(path)
+
+const isReactClass = path => {
+  return isClass(path) && hasSuperClass(path) && hasValidRenderMethod(getRenderMethod(path))
+}
+
+export default isReactClass

src/util/isReactFunction.js

@@ -0,0 +1,12 @@
+import * as t from 'babel-types'
+
+import getBody from './getBody'
+import hasReturnStatement from './hasReturnStatement'
+
+const isFunction = path => {
+  return t.isArrowFunctionExpression(path) || t.isFunctionDeclaration(path) || t.isFunctionExpression(path)
+}
+
+const isReactFunction = path => isFunction(path) && hasReturnStatement(getBody(path))
+
+export default isReactFunction

src/util/isReactImport.js

@@ -0,0 +1,5 @@
+import * as t from 'babel-types'
+
+const isReactImport = ({ node: { source } }) => t.isStringLiteral(source, { value: 'react' })
+
+export default isReactImport

src/util/isValidExpression.js

@@ -0,0 +1,13 @@
+import * as t from 'babel-types'
+import _ from 'lodash'
+
+const isValidExpression = ({ node: { left } }, names) => {
+  if (!t.isMemberExpression(left)) return false
+
+  const { property } = left
+  const { name } = property
+
+  return t.isIdentifier(property) && _.includes(names, name)
+}
+
+export default isValidExpression

src/util/isValidProperty.js

@@ -0,0 +1,11 @@
+import * as t from 'babel-types'
+import _ from 'lodash'
+
+const isValidProperty = (path, names) => {
+  const { node: { key } } = path
+  const { name } = key
+
+  return t.isIdentifier(key) && _.includes(names, name)
+}
+
+export default isValidProperty

src/util/types.js

@@ -0,0 +1,7 @@
+import * as t from 'babel-types'
+
+export const isArrayValue = path => t.isArrayExpression(path)
+
+export const isObjectValue = path => t.isObjectExpression(path)
+
+export const isStaticProperty = ({ node }) => !!node.static

src/visitors/entryVisitor.js

@@ -0,0 +1,14 @@
+import {
+  getEntryIdentifier,
+  isReactClass,
+  isReactFunction,
+} from '../util'
+
+const entriesVisitor = {
+  'Class|Function'(path, state) {
+    if (!isReactClass(path) && !isReactFunction(path)) return
+    state.createEntry(getEntryIdentifier(path), path)
+  },
+}
+
+export default entriesVisitor

src/visitors/importVisitor.js

@@ -0,0 +1,12 @@
+import { isReactImport } from '../util'
+
+const importVisitor = {
+  ImportDeclaration(path, state) {
+    if (!isReactImport(path)) return
+
+    state.hasImport = true
+    path.stop()
+  },
+}
+
+export default importVisitor

src/visitors/index.js

@@ -0,0 +1,3 @@
+export { default as entryVisitor } from './entryVisitor'
+export { default as importVisitor } from './importVisitor'
+export { default as propVisitor } from './propVisitor'