From 062409fa4b37f7ccfd2b216936d7fafdf0c6cf07 Mon Sep 17 00:00:00 2001
From: Ryan Kennedy <ryaninvents@users.noreply.github.com>
Date: Wed, 12 Jun 2019 16:06:11 -0400
Subject: [PATCH] feat: simplify custom context handling

---
 src/hooks/useReduxContext.js |  5 ++---
 src/hooks/useSelector.js     | 17 ++++++++++++++---
 src/hooks/useStore.js        |  9 +++++++--
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/src/hooks/useReduxContext.js b/src/hooks/useReduxContext.js
index 2a526ac2e..903a11c44 100644
--- a/src/hooks/useReduxContext.js
+++ b/src/hooks/useReduxContext.js
@@ -6,7 +6,6 @@ import { ReactReduxContext } from '../components/Context'
  * A hook to access the value of the `ReactReduxContext`. This is a low-level
  * hook that you should usually not need to call directly.
  *
- * @param {Function} [context=ReactReduxContext] Context passed to your `<Provider>`, if you're not using the default.
  * @returns {any} the value of the `ReactReduxContext`
  *
  * @example
@@ -19,8 +18,8 @@ import { ReactReduxContext } from '../components/Context'
  *   return <div>{store.getState()}</div>
  * }
  */
-export function useReduxContext(context = ReactReduxContext) {
-  const contextValue = useContext(context)
+export function useReduxContext() {
+  const contextValue = useContext(ReactReduxContext)
 
   invariant(
     contextValue,
diff --git a/src/hooks/useSelector.js b/src/hooks/useSelector.js
index 9d272ee61..2b88b16d1 100644
--- a/src/hooks/useSelector.js
+++ b/src/hooks/useSelector.js
@@ -1,6 +1,13 @@
-import { useReducer, useRef, useEffect, useMemo, useLayoutEffect } from 'react'
+import {
+  useReducer,
+  useRef,
+  useEffect,
+  useMemo,
+  useLayoutEffect,
+  useContext
+} from 'react'
 import invariant from 'invariant'
-import { useReduxContext } from './useReduxContext'
+import { useReduxContext as useDefaultReduxContext } from './useReduxContext'
 import Subscription from '../utils/Subscription'
 import { ReactReduxContext } from '../components/Context'
 
@@ -100,10 +107,14 @@ function useSelectorWithStoreAndSubscription(
  * @returns {Function} A `useSelector` hook bound to the specified context.
  */
 export function createSelectorHook(context = ReactReduxContext) {
+  const useReduxContext =
+    context === ReactReduxContext
+      ? useDefaultReduxContext
+      : () => useContext(context)
   return function useSelector(selector, equalityFn = refEquality) {
     invariant(selector, `You must pass a selector to useSelectors`)
 
-    const { store, subscription: contextSub } = useReduxContext(context)
+    const { store, subscription: contextSub } = useReduxContext()
 
     return useSelectorWithStoreAndSubscription(
       selector,
diff --git a/src/hooks/useStore.js b/src/hooks/useStore.js
index 376578fdb..82edaed01 100644
--- a/src/hooks/useStore.js
+++ b/src/hooks/useStore.js
@@ -1,5 +1,6 @@
+import { useContext } from 'react'
 import { ReactReduxContext } from '../components/Context'
-import { useReduxContext } from './useReduxContext'
+import { useReduxContext as useDefaultReduxContext } from './useReduxContext'
 
 /**
  * Hook factory, which creates a `useStore` hook bound to a given context.
@@ -8,8 +9,12 @@ import { useReduxContext } from './useReduxContext'
  * @returns {Function} A `useStore` hook bound to the specified context.
  */
 export function createStoreHook(context = ReactReduxContext) {
+  const useReduxContext =
+    context === ReactReduxContext
+      ? useDefaultReduxContext
+      : () => useContext(context).store
   return function useStore() {
-    const { store } = useReduxContext(context)
+    const { store } = useReduxContext()
     return store
   }
 }