Another pass at overloading map*ToProps

Author: tgriesser

src/components/connect.js

@@ -17,14 +17,21 @@ const defaultMergeProps = (stateProps, dispatchProps, parentProps) => ({
 function isFunction(fn) {
   return typeof fn === 'function'
 }
-function isFactory(fn) {
-  return isFunction(fn) && fn.length === 0 && isFunction(fn())
-}
 
 function getDisplayName(WrappedComponent) {
   return WrappedComponent.displayName || WrappedComponent.name || 'Component'
 }
 
+function checkStateShape(stateProps, dispatch) {
+  invariant(
+    isPlainObject(stateProps),
+    '`%sToProps` must return an object. Instead received %s.',
+    dispatch ? 'mapDispatch' : 'mapState',
+    stateProps
+  )
+  return stateProps
+}
+
 // Helps track hot reloading.
 let nextVersion = 0
 
@@ -35,9 +42,6 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
     wrapActionCreators(mapDispatchToProps) :
     mapDispatchToProps || defaultMapDispatchToProps
 
-  const mapStateFactory = isFactory(mapState) ? mapState : (() => mapState)
-  const mapDispatchFactory = isFactory(mapDispatch) ? mapDispatch : (() => mapDispatch)
-
   const finalMergeProps = mergeProps || defaultMergeProps
   const { pure = true, withRef = false } = options
 
@@ -62,45 +66,48 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
           `or explicitly pass "store" as a prop to "${this.constructor.displayName}".`
         )
 
-        this.configure()
         const storeState = this.store.getState()
         this.state = { storeState }
         this.clearCache()
       }
 
-      configure() {
-        this.finalMapStateToProps = mapStateFactory()
-        this.finalMapDispatchToProps = mapDispatchFactory()
-        this.doStatePropsDependOnOwnProps = this.finalMapStateToProps.length !== 1
-        this.doDispatchPropsDependOnOwnProps = this.finalMapDispatchToProps.length !== 1
-      }
-
       computeStateProps(store, props) {
+        if (!this.finalMapStateToProps) {
+          return this.configureFinalMapState(store, props)
+        }
         const state = store.getState()
         const stateProps = this.doStatePropsDependOnOwnProps ?
           this.finalMapStateToProps(state, props) :
           this.finalMapStateToProps(state)
 
-        invariant(
-          isPlainObject(stateProps),
-          '`mapStateToProps` must return an object. Instead received %s.',
-          stateProps
-        )
-        return stateProps
+        return checkStateShape(stateProps)
+      }
+
+      configureFinalMapState(store, props) {
+        const mappedState = mapState(store.getState(), props)
+        const isFactory = isFunction(mappedState)
+        this.finalMapStateToProps = isFactory ? mappedState : mapState
+        this.doStatePropsDependOnOwnProps = this.finalMapStateToProps.length !== 1
+        return isFactory ? this.computeStateProps(store, props) : checkStateShape(mappedState)
       }
 
       computeDispatchProps(store, props) {
+        if (!this.finalMapDispatchToProps) {
+          return this.configureFinalMapDispatch(store, props)
+        }
         const { dispatch } = store
         const dispatchProps = this.doDispatchPropsDependOnOwnProps ?
           this.finalMapDispatchToProps(dispatch, props) :
           this.finalMapDispatchToProps(dispatch)
+        return checkStateShape(dispatchProps, true)
+      }
 
-        invariant(
-          isPlainObject(dispatchProps),
-          '`mapDispatchToProps` must return an object. Instead received %s.',
-          dispatchProps
-        )
-        return dispatchProps
+      configureFinalMapDispatch(store, props) {
+        const mappedDispatch = mapDispatch(store.dispatch, props)
+        const isFactory = isFunction(mappedDispatch)
+        this.finalMapDispatchToProps = isFactory ? mappedDispatch : mapDispatch
+        this.doDispatchPropsDependOnOwnProps = this.finalMapDispatchToProps.length !== 1
+        return isFactory ? this.computeDispatchProps(store, props) : checkStateShape(mappedDispatch, true)
       }
 
       computeMergedProps(stateProps, dispatchProps, parentProps) {
@@ -181,6 +188,8 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
         this.haveOwnPropsChanged = true
         this.hasStoreStateChanged = true
         this.renderedElement = null
+        this.finalMapDispatchToProps = null
+        this.finalMapStateToProps = null
       }
 
       handleChange() {
@@ -282,7 +291,6 @@ export default function connect(mapStateToProps, mapDispatchToProps, mergeProps,
 
         // We are hot reloading!
         this.version = version
-        this.configure()
         this.trySubscribe()
         this.clearCache()
       }