ronantakizawa/codereview-bench huggingface dataset:
Repository: pmndrs/zustand
File: tests/persistSync.test.tsx (TypeScript)

Code before:
        expect(e?.message).toBe('getItem error')
      },
    })
  )

  expect(postRehydrationCallbackCallCount).toBe(1)
})

it('can persist state', () => {
  let setItemCallCount = 0

  const storage = {
    getItem: () => null,
    setItem: (name: string, value: string) => {
      setItemCallCount++
      expect(name).toBe('test-storage')
      expect(value).toBe(
        JSON.stringify({
          state: { count: 42 },
          version: 0,
        })
      )
    },
  }

  const useStore = create<any>(
    persist(() => ({ count: 0 }), {
      name: 'test-storage',
      getStorage: () => storage,
      onRehydrateStorage: () => (_, error) => {
        expect(error).toBeUndefined()
      },
    })
  )

  expect(useStore.getState()).toEqual({ count: 0 })
  useStore.setState({ count: 42 })
  expect(useStore.getState()).toEqual({ count: 42 })
  expect(setItemCallCount).toBe(1)
})

it('can migrate persisted state', () => {
  let migrateCallCount = 0
  let setItemCallCount = 0

  const storage = {
    getItem: () =>
      JSON.stringify({
        state: { count: 42 },
        version: 12,
      }),

Review comment (suggestion):
Maybe using real interfaces instead of `any` could help to test the typing integration? (I'm speaking of every test in this file, not this `create` specifically)

Code after:
          postRehydrationCallbackCallCount++
          expect(e?.message).toBe('getItem error')
        },
      })
    )

    expect(postRehydrationCallbackCallCount).toBe(1)
  })

  it('can persist state', () => {
    let setItemCallCount = 0

    const storage = {
      getItem: () => null,
      setItem: (name: string, value: string) => {
        setItemCallCount++
        expect(name).toBe('test-storage')
        expect(value).toBe(
          JSON.stringify({
            state: { count: 42 },
            version: 0,
          })
        )
      },
    }

    const useStore = create(
      persist(() => ({ count: 0 }), {
        name: 'test-storage',
        getStorage: () => storage,
        onRehydrateStorage: () => (_, error) => {
          expect(error).toBeUndefined()
        },
      })
    )

    expect(useStore.getState()).toEqual({ count: 0 })
    useStore.setState({ count: 42 })
    expect(useStore.getState()).toEqual({ count: 42 })
    expect(setItemCallCount).toBe(1)
  })

  it('can migrate persisted state', () => {
    let migrateCallCount = 0
    let setItemCallCount = 0

    const storage = {
      getItem: () =>
        JSON.stringify({
          state: { count: 42 },
          version: 12,
Repository: pmndrs/zustand
File: src/middleware.ts (TypeScript)

Code before:
  } = options || {}

  let storage: StateStorage | undefined

  try {
    storage = getStorage()
  } catch (e) {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
  }

  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
        set(...args)
      },
      get,
      api
    )
  }

  const thenableSerialize = toThenable(serialize)

  const setItem = (): Thenable<void> | void => {
    const state = { ...get() }

    if (whitelist) {
      ;(Object.keys(state) as (keyof S)[]).forEach((key) => {
        !whitelist.includes(key) && delete state[key]
      })
    }
    if (blacklist) {
      blacklist.forEach((key) => delete state[key])
    }

    return thenableSerialize({ state, version }).then((serializedValue) => {
      if (storage) {
        storage.setItem(name, serializedValue)
      } else {
        console.warn(
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
      }
    })
  }

  const savedSetState = api.setState

  api.setState = (state, replace) => {

Review comment (suggestion):
Does it return`void`? ~I'm fine with inferring return type (= no explicit return type), because we don't use it anywhere.~

Code after:
    name,
    getStorage = () => localStorage,
    serialize = JSON.stringify as (state: StorageValue<S>) => string,
    deserialize = JSON.parse as (str: string) => StorageValue<S>,
    blacklist,
    whitelist,
    onRehydrateStorage,
    version = 0,
    migrate,
  } = options || {}

  let storage: StateStorage | undefined

  try {
    storage = getStorage()
  } catch (e) {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
  }

  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
        set(...args)
      },
      get,
      api
    )
  }

  const thenableSerialize = toThenable(serialize)

  const setItem = (): Thenable<void> => {
    const state = { ...get() }

    if (whitelist) {
      ;(Object.keys(state) as (keyof S)[]).forEach((key) => {
        !whitelist.includes(key) && delete state[key]
      })
    }
    if (blacklist) {
      blacklist.forEach((key) => delete state[key])
    }

    let errorInSync: Error | undefined
    const thenable = thenableSerialize({ state, version })
      .then((serializedValue) =>
        (storage as StateStorage).setItem(name, serializedValue)
      )
Repository: pmndrs/zustand
File: src/middleware.ts (TypeScript)

Code before:
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
        set(...args)
      },
      get,
      api
    )
  }

  const thenableSerialize = toThenable(serialize)

  const setItem = (): Thenable<void> | void => {
    const state = { ...get() }

    if (whitelist) {
      ;(Object.keys(state) as (keyof S)[]).forEach((key) => {
        !whitelist.includes(key) && delete state[key]
      })
    }
    if (blacklist) {
      blacklist.forEach((key) => delete state[key])
    }

    return thenableSerialize({ state, version }).then((serializedValue) => {
      if (storage) {
        storage.setItem(name, serializedValue)
      } else {
        console.warn(
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
      }
    })
  }

  const savedSetState = api.setState

  api.setState = (state, replace) => {
    savedSetState(state, replace)
    void setItem()
  }

  // rehydrate initial state with existing stored state

  // a workaround to solve the issue of not storing rehydrated state in sync storage
  // the set(state) value would be later overridden with initial state by create()
  // to avoid this, we merge the state from localStorage into the initial state.
  // TODO: find a better solution for this
  let stateFromStorageInSync: S | undefined
  const postRehydrationCallback = onRehydrateStorage?.(get()) || undefined
  // bind is used to avoid `TypeError: Illegal invocation` error
  toThenable(storage.getItem.bind(storage))(name)

Review comment (suggestion):
As we don't return the `storage.setItem` result, `return` in L286 is useless.
We don't use the return value of setItem anywhere (or did I miss it?), could we remove `return` in L286? (and that will change other code, if it works.)

Code after:
    storage = getStorage()
  } catch (e) {
    // prevent error if the storage is not defined (e.g. when server side rendering a page)
  }

  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `Persist middleware: unable to update ${name}, the given storage is currently unavailable.`
        )
        set(...args)
      },
      get,
      api
    )
  }

  const thenableSerialize = toThenable(serialize)

  const setItem = (): Thenable<void> => {
    const state = { ...get() }

    if (whitelist) {
      ;(Object.keys(state) as (keyof S)[]).forEach((key) => {
        !whitelist.includes(key) && delete state[key]
      })
    }
    if (blacklist) {
      blacklist.forEach((key) => delete state[key])
    }

    let errorInSync: Error | undefined
    const thenable = thenableSerialize({ state, version })
      .then((serializedValue) =>
        (storage as StateStorage).setItem(name, serializedValue)
      )
      .catch((e) => {
        errorInSync = e
      })
    if (errorInSync) {
      throw errorInSync
    }
    return thenable
  }

  const savedSetState = api.setState

  api.setState = (state, replace) => {
    savedSetState(state, replace)
    void setItem()
Repository: pmndrs/zustand
File: tests/persistAsync.test.tsx (TypeScript)

Code before:

describe('persist middleware with async configuration', () => {
  it('can rehydrate state', async () => {
    const storage = {
      getItem: async (name: string) =>
        JSON.stringify({
          state: { count: 42, name },
          version: 0,
        }),
      setItem: () => {},
    }

    const spy = jest.fn()
    const useStore = create(
      persist(
        () => ({
          count: 0,
          name: 'empty',
        }),
        {
          name: 'test-storage',
          getStorage: () => storage,
          onRehydrateStorage: () => spy,
        }
      )
    )

    function Counter() {
      const { count, name } = useStore()
      return (
        <div>
          count: {count}, name: {name}
        </div>
      )
    }

    const { findByText } = render(<Counter />)

    await findByText('count: 0, name: empty')
    await findByText('count: 42, name: test-storage')
    expect(spy).toBeCalledWith({ count: 42, name: 'test-storage' }, undefined)
  })

  it('can throw rehydrate error', async () => {
    let postRehydrationCallbackCallCount = 0

    const storage = {
      getItem: async () => {
        throw new Error('getItem error')
      },
      setItem: () => {},

Review comment (suggestion):
As written in the description, those `expect()` calls never worked. They resulted in an error being thrown [here](https://github.com/pmndrs/zustand/blob/master/src/middleware.ts#L363) which only a second `catch()` would catch (I actually tried).

So I'd even go so far as to say that this way of testing callbacks is an anti pattern that should be avoided.

Code after:

describe('persist middleware with async configuration', () => {
  const consoleError = console.error
  afterEach(() => {
    console.error = consoleError
  })

  it('can rehydrate state', async () => {
    const onRehydrateStorageSpy = jest.fn()
    const storage = {
      getItem: async (name: string) =>
        JSON.stringify({
          state: { count: 42, name },
          version: 0,
        }),
      setItem: () => {},
    }

    const useStore = create(
      persist(
        () => ({
          count: 0,
          name: 'empty',
        }),
        {
          name: 'test-storage',
          getStorage: () => storage,
          onRehydrateStorage: () => onRehydrateStorageSpy,
        }
      )
    )

    function Counter() {
      const { count, name } = useStore()
      return (
        <div>
          count: {count}, name: {name}
        </div>
      )
    }

    const { findByText } = render(<Counter />)

    await findByText('count: 0, name: empty')
    await findByText('count: 42, name: test-storage')
    expect(onRehydrateStorageSpy).toBeCalledWith(
      { count: 42, name: 'test-storage' },
      undefined
    )
  })
Repository: pmndrs/zustand
File: src/context.ts (TypeScript)

Code before:
import {
  ReactNode,
  createElement,
  createContext as reactCreateContext,
  useContext,
  useMemo,
  useRef,
} from 'react'
import { EqualityChecker, UseStore } from 'zustand'
import { State, StateSelector } from './vanilla'

export interface UseStoreData<T extends State> {
  (): T
  <U>(selector: StateSelector<T, U>, equalityFn?: EqualityChecker<U>): U
}

function createContext<TState extends State>() {
  const ZustandContext = reactCreateContext<UseStore<TState> | undefined>(
    undefined
  )

  const Provider = ({
    // @ts-expect-error deprecated
    initialStore,
    createStore,
    children,
  }: {
    createStore: () => UseStore<TState>
    children: ReactNode
  }) => {
    const storeRef = useRef<UseStore<TState>>()

    if (!storeRef.current) {
      if (initialStore) {
        console.warn(
          'Provider initialStore is deprecated and will be removed in the next version.'
        )

Review comment (suggestion):
```suggestion
interface UseStoreData<T extends State> {
```
Yeah, let remove `export` for this time, because I'm not confident with the name.

Code after:
import {
  ReactNode,
  createElement,
  createContext as reactCreateContext,
  useContext,
  useMemo,
  useRef,
} from 'react'
import { EqualityChecker, UseStore } from 'zustand'
import { State, StateSelector } from './vanilla'

interface UseStoreData<T extends State> {
  (): T
  <U>(selector: StateSelector<T, U>, equalityFn?: EqualityChecker<U>): U
}

function createContext<TState extends State>() {
  const ZustandContext = reactCreateContext<UseStore<TState> | undefined>(
    undefined
  )

  const Provider = ({
    // @ts-expect-error deprecated
    initialStore,
    createStore,
    children,
  }: {
    createStore: () => UseStore<TState>
    children: ReactNode
  }) => {
    const storeRef = useRef<UseStore<TState>>()

    if (!storeRef.current) {
      if (initialStore) {
        console.warn(
          'Provider initialStore is deprecated and will be removed in the next version.'
        )
Repository: pmndrs/zustand
File: src/middleware.ts (TypeScript)

Code before:
      const result = fn(input)
      if (result instanceof Promise) {
        return result as Thenable<Result>
      }
      return {
        then(onFulfilled) {
          return toThenable(onFulfilled)(result as Result)
        },
        catch(_onRejected) {
          return this as Thenable<any>
        },
      }
    } catch (e: any) {
      return {
        then(_onFulfilled) {
          return this as Thenable<any>
        },
        catch(onRejected) {
          return toThenable(onRejected)(e)
        },
      }
    }
  }

export const persist =
  <S extends State>(config: StateCreator<S>, _options: PersistOptions<S>) =>
  (set: SetState<S>, get: GetState<S>, api: PersistApi<S>): S => {
    let options = {
      getStorage: () => localStorage,
      serialize: JSON.stringify as (state: StorageValue<S>) => string,
      deserialize: JSON.parse as (str: string) => StorageValue<Partial<S>>,
      partialize: (state: S) => state,
      version: 0,
      merge: (persistedState: any, currentState: S) => ({
        ...currentState,
        ...persistedState,
      }),
      ..._options,
    }

    if (options.blacklist || options.whitelist) {
      console.warn(
        `The ${
          options.blacklist ? 'blacklist' : 'whitelist'
        } option is deprecated and will be removed in the next version. Please use the 'partialize' option instead.`
      )
    }

    let hasHydrated = false
    const onHydrateCallbacks: Parameters<
      PersistApi<S>['persist']['onHydrate']

Review comment (suggestion):
Not 100% sure about the use case, but we want `CustomSetState`, `CustomGetState`, `CustomStoreApi`, like some other middleware.

Code after:
      const result = fn(input)
      if (result instanceof Promise) {
        return result as Thenable<Result>
      }
      return {
        then(onFulfilled) {
          return toThenable(onFulfilled)(result as Result)
        },
        catch(_onRejected) {
          return this as Thenable<any>
        },
      }
    } catch (e: any) {
      return {
        then(_onFulfilled) {
          return this as Thenable<any>
        },
        catch(onRejected) {
          return toThenable(onRejected)(e)
        },
      }
    }
  }

export const persist =
  <
    S extends State,
    CustomSetState extends SetState<S>,
    CustomGetState extends GetState<S>,
    CustomStoreApi extends StoreApiWithPersist<S>
  >(
    config: (
      set: CustomSetState,
      get: CustomGetState,
      api: CustomStoreApi
    ) => S,
    baseOptions: PersistOptions<S>
  ) =>
  (set: CustomSetState, get: CustomGetState, api: CustomStoreApi): S => {
    let options = {
      getStorage: () => localStorage,
      serialize: JSON.stringify as (state: StorageValue<S>) => string,
      deserialize: JSON.parse as (str: string) => StorageValue<Partial<S>>,
      partialize: (state: S) => state,
      version: 0,
      merge: (persistedState: any, currentState: S) => ({
        ...currentState,
        ...persistedState,
      }),
      ...baseOptions,
    }
Repository: pmndrs/zustand
File: tests/middlewareTypes.test.tsx (TypeScript)

Code before:
      return <></>
    }
    TestComponent
  })

  it('combine', () => {
    const useStore = create(
      combine({ count: 1 }, (set, get) => ({
        inc: () => set({ count: get().count + 1 }, false),
      }))
    )
    const TestComponent = () => {
      useStore((s) => s.count) * 2
      useStore((s) => s.inc)()
      useStore().count * 2
      useStore().inc()
      useStore.getState().count * 2
      useStore.getState().inc()
      return <></>
    }
    TestComponent
  })

  it('persist', () => {
    const useStore = create(
      persist<CounterState>(
        (set, get) => ({
          count: 1,
          inc: () => set({ count: get().count + 1 }, false),
        }),
        { name: 'prefix' }
      )
    )
    const TestComponent = () => {
      useStore((s) => s.count) * 2
      useStore((s) => s.inc)()
      useStore().count * 2
      useStore().inc()
      useStore.getState().count * 2
      useStore.getState().inc()
      return <></>
    }
    TestComponent
  })
})

describe('counter state spec (double middleware)', () => {
  it('devtools & immer', () => {
    const useStore = create<CounterState>(
      devtools(
        immer((set, get) => ({

Review comment (suggestion):
We would like to keep type annotation in `create`.

Code after:
        { name: 'prefix' }
      )
    )
    const TestComponent = () => {
      useStore((s) => s.count) * 2
      useStore((s) => s.inc)()
      useStore().count * 2
      useStore().inc()
      useStore.getState().count * 2
      useStore.getState().inc()
      useStore.subscribe(
        (state) => state.count,
        (count) => console.log(count * 2)
      )
      return <></>
    }
    TestComponent
  })

  it('combine', () => {
    const useStore = create(
      combine({ count: 1 }, (set, get) => ({
        inc: () => set({ count: get().count + 1 }, false),
      }))
    )
    const TestComponent = () => {
      useStore((s) => s.count) * 2
      useStore((s) => s.inc)()
      useStore().count * 2
      useStore().inc()
      useStore.getState().count * 2
      useStore.getState().inc()
      return <></>
    }
    TestComponent
  })

  it('persist', () => {
    const useStore = create<
      CounterState,
      SetState<CounterState>,
      GetState<CounterState>,
      StoreApiWithPersist<CounterState>
    >(
      persist(
        (set, get) => ({
          count: 1,
          inc: () => set({ count: get().count + 1 }, false),
        }),
        { name: 'prefix' }
      )
Repository: pmndrs/zustand
File: src/middleware.ts (TypeScript)

Code before:
        })
    }

    api.persist = {
      setOptions: (newOptions) => {
        options = {
          ...options,
          ...newOptions,
        }

        if (newOptions.getStorage) {
          storage = newOptions.getStorage()
        }
      },
      clearStorage: () => {
        storage?.removeItem?.(options.name)
      },
      rehydrate: () => hydrate() as Promise<void>,
      hasHydrated: () => hasHydrated,
      onHydrate: (cb) => {
        onHydrateCallbacks.push(cb)
      },
      onFinishHydration: (cb) => {
        onFinishHydrationCallbacks.push(cb)
      },
      useHydration: () => {
        const [hasHydrated, setHasHydrated] = useState(
          api.persist.hasHydrated()
        )

        useEffect(() => {
          api.persist.onHydrate(() => {
            setHasHydrated(false)
          })

          api.persist.onFinishHydration(() => {
            setHasHydrated(true)
          })
        }, [])

        return {
          hasHydrated,
        }
      },
    }

    hydrate()

    return stateFromStorage || configResult
  }


Review comment (suggestion):
I don't think it's a good idea for `zustand/middleware` to depend on `react`. All middlewares including persist should be able to be used in vanilla.

So, I would suggest to not have it in persist api, but in docs as a reference implementation.
Also, if we use a hook, useEffect should be able to unsubscribe the callbacks.

```js
const useStore = create(persist(...))
const useHydration = () => {
  const [hydrated, setHydrated] = useState(useStore.persist.hasHydrated)
  useEffect(() => {
    const unsub1 = useStore.persist.onHydrate(() => setHydrated(false)) // not sure this is required in practice
    const unsub2 = useStore.persist.onFinishHydration(() => setHydrated(true))
    setHydrated(useStore.persist.hasHydrated())
    return () => {
      unsub1()
      unsub2()
    }
  }, [])
  return hydrated
}
```

Code after:
          ...newOptions,
        }

        if (newOptions.getStorage) {
          storage = newOptions.getStorage()
        }
      },
      clearStorage: () => {
        storage?.removeItem?.(options.name)
      },
      rehydrate: () => hydrate() as Promise<void>,
      hasHydrated: () => hasHydrated,
      onHydrate: (cb) => {
        hydrationListeners.add(cb)

        return () => {
          hydrationListeners.delete(cb)
        }
      },
      onFinishHydration: (cb) => {
        finishHydrationListeners.add(cb)

        return () => {
          finishHydrationListeners.delete(cb)
        }
      },
    }

    hydrate()

    return stateFromStorage || configResult
  }