diff --git a/src/Store.js b/src/Store.js
deleted file mode 100644
index 3035d9d..0000000
--- a/src/Store.js
+++ /dev/null
@@ -1,14 +0,0 @@
-import { has, uniq } from 'lodash'
-
-export default class Store {
-  props = {}
-
-  add = (name, property) => {
-    if (!this.has(name)) this.props[name] = []
-    this.props[name].push(property)
-  }
-
-  has = name => has(this.props, name)
-
-  get = name => uniq(this.props[name]).sort()
-}
diff --git a/src/assertions.js b/src/assertions.js
index e695610..856ecde 100644
--- a/src/assertions.js
+++ b/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
 }
diff --git a/src/helpers.js b/src/helpers.js
deleted file mode 100644
index 761e3da..0000000
--- a/src/helpers.js
+++ /dev/null
@@ -1,14 +0,0 @@
-import * as t from 'babel-types'
-
-export const findClassIdentifier = (path) => {
-  const declaration = path.findParent(parentPath => parentPath.isClassDeclaration())
-  return declaration.node.id.name
-}
-
-export const generateExpression = (store, identifier) => {
-  const props = store.get(identifier).map(prop => t.stringLiteral(prop))
-  const left = t.memberExpression(t.identifier(identifier), t.identifier('handledProps'))
-  const expression = t.assignmentExpression('=', left, t.arrayExpression(props))
-
-  return t.expressionStatement(expression)
-}
diff --git a/src/index.js b/src/index.js
index a603f15..cb72dcf 100644
--- a/src/index.js
+++ b/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
diff --git a/src/util/Store.js b/src/util/Store.js
new file mode 100644
index 0000000..80220dd
--- /dev/null
+++ b/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)
+}
diff --git a/src/util/createPropertyExpression.js b/src/util/createPropertyExpression.js
new file mode 100644
index 0000000..502e281
--- /dev/null
+++ b/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
diff --git a/src/util/getBody.js b/src/util/getBody.js
new file mode 100644
index 0000000..1bac6ea
--- /dev/null
+++ b/src/util/getBody.js
@@ -0,0 +1,3 @@
+const getBody = ({ node: { body } }) => body.body
+
+export default getBody
diff --git a/src/util/getClassDeclaration.js b/src/util/getClassDeclaration.js
new file mode 100644
index 0000000..a78a887
--- /dev/null
+++ b/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
diff --git a/src/util/getEntryIdentifier.js b/src/util/getEntryIdentifier.js
new file mode 100644
index 0000000..dfb8ced
--- /dev/null
+++ b/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
diff --git a/src/util/getExpressionIdentifier.js b/src/util/getExpressionIdentifier.js
new file mode 100644
index 0000000..e994d1e
--- /dev/null
+++ b/src/util/getExpressionIdentifier.js
@@ -0,0 +1,3 @@
+const getExpressionIdentifier = ({ node: { left: { object } } }) => object.name
+
+export default getExpressionIdentifier
diff --git a/src/util/hasReturnStatement.js b/src/util/hasReturnStatement.js
new file mode 100644
index 0000000..896f0a0
--- /dev/null
+++ b/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
diff --git a/src/util/index.js b/src/util/index.js
new file mode 100644
index 0000000..aef0896
--- /dev/null
+++ b/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'
diff --git a/src/util/insertAfterPath.js b/src/util/insertAfterPath.js
new file mode 100644
index 0000000..17dc4a2
--- /dev/null
+++ b/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
diff --git a/src/util/isReactClass.js b/src/util/isReactClass.js
new file mode 100644
index 0000000..237ffc0
--- /dev/null
+++ b/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
diff --git a/src/util/isReactFunction.js b/src/util/isReactFunction.js
new file mode 100644
index 0000000..bc29208
--- /dev/null
+++ b/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
diff --git a/src/util/isReactImport.js b/src/util/isReactImport.js
new file mode 100644
index 0000000..e0f533e
--- /dev/null
+++ b/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
diff --git a/src/util/isValidExpression.js b/src/util/isValidExpression.js
new file mode 100644
index 0000000..c1ec62c
--- /dev/null
+++ b/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
diff --git a/src/util/isValidProperty.js b/src/util/isValidProperty.js
new file mode 100644
index 0000000..2ef5d7e
--- /dev/null
+++ b/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
diff --git a/src/util/types.js b/src/util/types.js
new file mode 100644
index 0000000..05ff08e
--- /dev/null
+++ b/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
diff --git a/src/visitors/entryVisitor.js b/src/visitors/entryVisitor.js
new file mode 100644
index 0000000..f14d1b8
--- /dev/null
+++ b/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
diff --git a/src/visitors/importVisitor.js b/src/visitors/importVisitor.js
new file mode 100644
index 0000000..5121245
--- /dev/null
+++ b/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
diff --git a/src/visitors/index.js b/src/visitors/index.js
new file mode 100644
index 0000000..d0d2038
--- /dev/null
+++ b/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'
diff --git a/src/visitors/propVisitor.js b/src/visitors/propVisitor.js
new file mode 100644
index 0000000..ca0fb89
--- /dev/null
+++ b/src/visitors/propVisitor.js
@@ -0,0 +1,54 @@
+import _ from 'lodash'
+import {
+  getClassDeclaration,
+  getEntryIdentifier,
+  getExpressionIdentifier,
+  isArrayValue,
+  isObjectValue,
+  isStaticProperty,
+  isValidExpression,
+  isValidProperty,
+} from '../util'
+
+const getArrayItems = ({ elements }) => _.map(elements, ({ value }) => value)
+
+const getObjectKeys = ({ properties }) => _.map(properties, ({ key: { name } }) => name)
+
+const propVisitor = {
+  AssignmentExpression(path, state) {
+    const identifier = getExpressionIdentifier(path)
+    const { node: { right } } = path
+
+    if (!state.hasEntry(identifier)) return
+
+    if (isValidExpression(path, ['handledProps']) && isArrayValue(right)) {
+      state.addProps(identifier, getArrayItems(right))
+      path.remove()
+
+      return
+    }
+
+    if (isValidExpression(path, ['defaultProps', 'propTypes']) && isObjectValue(right)) {
+      state.addProps(identifier, getObjectKeys(right))
+    }
+  },
+  ClassProperty(path, state) {
+    const identifier = getEntryIdentifier(getClassDeclaration(path))
+    const { node: { value } } = path
+
+    if (!state.hasEntry(identifier) || !isStaticProperty(path)) return
+
+    if (isValidProperty(path, ['handledProps']) && isArrayValue(value)) {
+      state.addProps(identifier, getArrayItems(value))
+      path.remove()
+
+      return
+    }
+
+    if (isValidProperty(path, ['defaultProps', 'propTypes']) && isObjectValue(value)) {
+      state.addProps(identifier, getObjectKeys(value))
+    }
+  },
+}
+
+export default propVisitor
diff --git a/test/fixtures/skiped/actual.js b/test/fixtures/skiped/actual.js
new file mode 100644
index 0000000..c069b5d
--- /dev/null
+++ b/test/fixtures/skiped/actual.js
@@ -0,0 +1,12 @@
+function Example() {
+  return null;
+}
+Example.defaultProps = {
+  active: true
+};
+Example.propTypes = {
+  children: PropTypes.node,
+  className: PropTypes.string
+};
+
+export default Example;
diff --git a/test/fixtures/skiped/expected.js b/test/fixtures/skiped/expected.js
new file mode 100644
index 0000000..c069b5d
--- /dev/null
+++ b/test/fixtures/skiped/expected.js
@@ -0,0 +1,12 @@
+function Example() {
+  return null;
+}
+Example.defaultProps = {
+  active: true
+};
+Example.propTypes = {
+  children: PropTypes.node,
+  className: PropTypes.string
+};
+
+export default Example;
diff --git a/test/fixtures/stateless-arrow/actual.js b/test/fixtures/stateless-arrow/actual.js
new file mode 100644
index 0000000..e746906
--- /dev/null
+++ b/test/fixtures/stateless-arrow/actual.js
@@ -0,0 +1,14 @@
+import React, { PropTypes } from 'react';
+
+const Example = () => {
+  return null;
+};
+Example.defaultProps = {
+  active: true
+};
+Example.propTypes = {
+  children: PropTypes.node,
+  className: PropTypes.string
+};
+
+export default Example;
diff --git a/test/fixtures/stateless-arrow/expected.js b/test/fixtures/stateless-arrow/expected.js
new file mode 100644
index 0000000..761d685
--- /dev/null
+++ b/test/fixtures/stateless-arrow/expected.js
@@ -0,0 +1,16 @@
+import React, { PropTypes } from 'react';
+
+const Example = (_temp = () => {
+  var _temp;
+
+  return null;
+}, Example.handledProps = ['active', 'children', 'className'], _temp);
+Example.defaultProps = {
+  active: true
+};
+Example.propTypes = {
+  children: PropTypes.node,
+  className: PropTypes.string
+};
+
+export default Example;
diff --git a/test/fixtures/stateless-assignment/actual.js b/test/fixtures/stateless-assignment/actual.js
new file mode 100644
index 0000000..e1f2ba0
--- /dev/null
+++ b/test/fixtures/stateless-assignment/actual.js
@@ -0,0 +1,14 @@
+import React, { PropTypes } from 'react';
+
+const Example = function () {
+  return null;
+};
+Example.defaultProps = {
+  active: true
+};
+Example.propTypes = {
+  children: PropTypes.node,
+  className: PropTypes.string
+};
+
+export default Example;
diff --git a/test/fixtures/stateless-assignment/expected.js b/test/fixtures/stateless-assignment/expected.js
new file mode 100644
index 0000000..72526fd
--- /dev/null
+++ b/test/fixtures/stateless-assignment/expected.js
@@ -0,0 +1,16 @@
+import React, { PropTypes } from 'react';
+
+const Example = (_temp = function () {
+  var _temp;
+
+  return null;
+}, Example.handledProps = ['active', 'children', 'className'], _temp);
+Example.defaultProps = {
+  active: true
+};
+Example.propTypes = {
+  children: PropTypes.node,
+  className: PropTypes.string
+};
+
+export default Example;