feat(middleware): subscribeWithSelector middleware (pmndrs#603)

* feat(middleware): subscribeWithSelector middleware, deprecating the core feature

* update readme

* update readme

readme.md

@@ -177,19 +177,10 @@ const useStore = create(() => ({ paw: true, snout: true, fur: true }))
 const paw = useStore.getState().paw
 // Listening to all changes, fires synchronously on every change
 const unsub1 = useStore.subscribe(console.log)
-// Listening to selected changes, in this case when "paw" changes
-const unsub2 = useStore.subscribe(console.log, state => state.paw)
-// Subscribe also supports an optional equality function
-const unsub3 = useStore.subscribe(console.log, state => [state.paw, state.fur], shallow)
-// Subscribe also exposes the previous value
-const unsub4 = useStore.subscribe((paw, previousPaw) => console.log(paw, previousPaw), state => state.paw)
 // Updating state, will trigger listeners
 useStore.setState({ paw: false })
 // Unsubscribe listeners
 unsub1()
-unsub2()
-unsub3()
-unsub4()
 // Destroying the store (removing all listeners)
 useStore.destroy()
 
@@ -198,6 +189,30 @@ function Component() {
   const paw = useStore(state => state.paw)
 ```
 
+### Using subscribe with selector
+
+If you need to subscribe with selector,
+`subscribeWithSelector` middleware will help.
+
+With this middleware `subscribe` accepts an additional signature:
+```ts
+subscribe(selector, callback, options?: { equalityFn, fireImmediately }): Unsubscribe
+```
+
+```js
+import { subscribeWithSelector } from 'zustand/middleware'
+const useStore = create(subscribeWithSelector(() => ({ paw: true, snout: true, fur: true })))
+
+// Listening to selected changes, in this case when "paw" changes
+const unsub2 = useStore.subscribe(state => state.paw, console.log)
+// Subscribe also exposes the previous value
+const unsub3 = useStore.subscribe(state => state.paw, (paw, previousPaw) => console.log(paw, previousPaw))
+// Subscribe also supports an optional equality function
+const unsub4 = useStore.subscribe(state => [state.paw, state.fur], console.log, { equalityFn: shallow })
+// Subscribe and fire immediately
+const unsub5 = useStore.subscribe(state => state.paw, console.log, { fireImmediately: true })
+```
+
 ## Using zustand without React
 
 Zustands core can be imported and used without the React dependency. The only difference is that the create function does not return a hook, but the api utilities.

src/middleware.ts

@@ -1,10 +1,15 @@
 import {
+  EqualityChecker,
   GetState,
   PartialState,
   SetState,
   State,
   StateCreator,
+  StateListener,
+  StateSelector,
+  StateSliceListener,
   StoreApi,
+  Subscribe,
 } from './vanilla'
 
 const DEVTOOLS = Symbol()
@@ -192,6 +197,104 @@ export const devtools =
     return initialState
   }
 
+type SubscribeWithSelector<T extends State> = {
+  (listener: StateListener<T>): () => void
+  <StateSlice>(
+    selector: StateSelector<T, StateSlice>,
+    listener: StateSliceListener<StateSlice>,
+    options?: {
+      equalityFn?: EqualityChecker<StateSlice>
+      fireImmediately?: boolean
+    }
+  ): () => void
+}
+
+export function subscribeWithSelector<S extends State>(
+  fn: (
+    set: SetState<S>,
+    get: GetState<S>,
+    api: Omit<StoreApi<S>, 'subscribe'> & {
+      subscribe: SubscribeWithSelector<S>
+      subscribeWithSelectorEnabled: true
+    }
+  ) => S
+): (
+  set: SetState<S>,
+  get: GetState<S>,
+  api: Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+) => S
+
+export function subscribeWithSelector<
+  S extends State,
+  CustomSetState extends SetState<S>
+>(
+  fn: (
+    set: CustomSetState,
+    get: GetState<S>,
+    api: Omit<StoreApi<S>, 'subscribe'> & {
+      subscribe: SubscribeWithSelector<S>
+      subscribeWithSelectorEnabled: true
+    }
+  ) => S
+): (
+  set: CustomSetState,
+  get: GetState<S>,
+  api: Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+) => S
+
+export function subscribeWithSelector<
+  S extends State,
+  CustomSetState extends SetState<S>,
+  CustomGetState extends GetState<S>,
+  CustomStoreApi extends Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+>(
+  fn: (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S
+): (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S
+
+export function subscribeWithSelector<
+  S extends State,
+  CustomSetState extends SetState<S>,
+  CustomGetState extends GetState<S>,
+  CustomStoreApi extends Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+>(fn: (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S) {
+  return (set: CustomSetState, get: CustomGetState, api: CustomStoreApi): S => {
+    const origSubscribe = api.subscribe as Subscribe<S>
+    api.subscribe = ((selector: any, optListener: any, options: any) => {
+      let listener: StateListener<S> = selector // if no selector
+      if (optListener) {
+        const equalityFn = options?.equalityFn || Object.is
+        let currentSlice = selector(api.getState())
+        listener = (state) => {
+          const nextSlice = selector(state)
+          if (!equalityFn(currentSlice, nextSlice)) {
+            const previousSlice = currentSlice
+            optListener((currentSlice = nextSlice), previousSlice)
+          }
+        }
+        if (options?.fireImmediately) {
+          optListener(currentSlice, currentSlice)
+        }
+      }
+      return origSubscribe(listener)
+    }) as SubscribeWithSelector<S>
+    api.subscribeWithSelectorEnabled = true
+    const initialState = fn(set, get, api)
+    return initialState
+  }
+}
+
 type Combine<T, U> = Omit<T, keyof U> & U
 export const combine =
   <

src/vanilla.ts

@@ -16,6 +16,9 @@ export type StateListener<T> = (state: T, previousState: T) => void
 export type StateSliceListener<T> = (slice: T, previousSlice: T) => void
 export interface Subscribe<T extends State> {
   (listener: StateListener<T>): () => void
+  /**
+   * @deprecated Please use `subscribeWithSelector` middleware
+   */
   <StateSlice>(
     listener: StateSliceListener<StateSlice>,
     selector?: StateSelector<T, StateSlice>,
@@ -88,6 +91,7 @@ export default function create<
     selector: StateSelector<TState, StateSlice> = getState as any,
     equalityFn: EqualityChecker<StateSlice> = Object.is
   ) => {
+    console.warn('[DEPRECATED] Please use `subscribeWithSelector` middleware')
     let currentSlice: StateSlice = selector(state)
     function listenerToAdd() {
       const nextSlice = selector(state)

tests/context.test.tsx

@@ -1,7 +1,8 @@
 import { Component as ClassComponent, useEffect, useState } from 'react'
 import { render } from '@testing-library/react'
-import create from 'zustand'
+import create, { StoreApi } from 'zustand'
 import createContext from 'zustand/context'
+import { subscribeWithSelector } from 'zustand/middleware'
 
 const consoleError = console.error
 afterEach(() => {
@@ -63,22 +64,32 @@ it('uses context store with selectors', async () => {
 })
 
 it('uses context store api', async () => {
-  const { Provider, useStoreApi } = createContext<CounterState>()
+  const { Provider, useStoreApi: useStoreApiOrig } =
+    createContext<CounterState>()
 
   const createStore = () =>
-    create<CounterState>((set) => ({
-      count: 0,
-      inc: () => set((state) => ({ count: state.count + 1 })),
-    }))
+    create(
+      subscribeWithSelector<{ count: number; inc: () => void }>((set) => ({
+        count: 0,
+        inc: () => set((state) => ({ count: state.count + 1 })),
+      }))
+    )
+
+  const useStoreApi = useStoreApiOrig as () => Omit<
+    StoreApi<CounterState>,
+    'subscribe'
+  > & {
+    subscribe: ReturnType<typeof createStore>['subscribe']
+  }
 
   function Counter() {
     const storeApi = useStoreApi()
     const [count, setCount] = useState(0)
     useEffect(
       () =>
         storeApi.subscribe(
-          () => setCount(storeApi.getState().count),
-          (state) => state.count
+          (state) => state.count,
+          () => setCount(storeApi.getState().count)
         ),
       [storeApi]
     )

tests/middlewareTypes.test.tsx

@@ -1,7 +1,14 @@
 import { produce } from 'immer'
 import type { Draft } from 'immer'
 import create, { UseBoundStore } from 'zustand'
-import { NamedSet, combine, devtools, persist, redux } from 'zustand/middleware'
+import {
+  NamedSet,
+  combine,
+  devtools,
+  persist,
+  redux,
+  subscribeWithSelector,
+} from 'zustand/middleware'
 import { State, StateCreator } from 'zustand/vanilla'
 
 type TImmerConfigFn<T extends State> = (fn: (draft: Draft<T>) => void) => void
@@ -309,3 +316,87 @@ it('should combine devtools and combine', () => {
   }
   TestComponent
 })
+
+it('should combine subscribeWithSelector and combine', () => {
+  const useStore = create(
+    subscribeWithSelector(
+      combine({ count: 1 }, (set, get) => ({
+        inc: () => set({ count: get().count + 1 }, false, 'inc'),
+      }))
+    )
+  )
+
+  const TestComponent = (): JSX.Element => {
+    useStore().count
+    useStore().inc()
+    useStore.getState().count
+    useStore.getState().inc()
+    useStore.subscribe(
+      (state) => state.count,
+      (count) => console.log(count * 2)
+    )
+
+    return <></>
+  }
+  TestComponent
+})
+
+it('should combine devtools and subscribeWithSelector', () => {
+  const useStore = create(
+    devtools(
+      subscribeWithSelector<
+        {
+          count: number
+          inc: () => void
+        },
+        NamedSet<{
+          count: number
+          inc: () => void
+        }>
+      >((set, get) => ({
+        count: 1,
+        inc: () => set({ count: get().count + 1 }, false, 'inc'),
+      }))
+    )
+  )
+
+  const TestComponent = (): JSX.Element => {
+    useStore().count
+    useStore().inc()
+    useStore.getState().count
+    useStore.getState().inc()
+    useStore.subscribe(
+      (state) => state.count,
+      (count) => console.log(count * 2)
+    )
+
+    return <></>
+  }
+  TestComponent
+})
+
+it('should combine devtools, subscribeWithSelector and combine', () => {
+  const useStore = create(
+    devtools(
+      subscribeWithSelector(
+        combine({ count: 1 }, (set, get) => ({
+          inc: () => set({ count: get().count + 1 }, false, 'inc'),
+        }))
+      )
+    )
+  )
+
+  const TestComponent = (): JSX.Element => {
+    useStore().count
+    useStore().inc()
+    useStore.getState().count
+    useStore.getState().inc()
+    useStore.subscribe(
+      (state) => state.count,
+      (count) => console.log(count * 2)
+    )
+
+    return <></>
+  }
+  TestComponent
+})