Allow functional updates to state

src/use-localstorage.ts

@@ -5,7 +5,7 @@ import {
   LOCAL_STORAGE_CHANGE_EVENT_NAME,
 } from './local-storage-events';
 import { isBrowser } from './is-browser'
-import { storage  } from './storage'
+import { storage } from './storage'
 import { useEffect, useState, useCallback } from 'react';
 
 /**
@@ -22,8 +22,9 @@ function tryParse(value: string) {
   }
 }
 
-export type LocalStorageNullableReturnValue<TValue> = [TValue | null, (newValue: TValue | null) => void, () => void];
-export type LocalStorageReturnValue<TValue> = [TValue, (newValue: TValue | null) => void, () => void];
+export type LocalStorageSetStateValue<TValue> = TValue | ((prevState: TValue | null) => TValue)
+export type LocalStorageNullableReturnValue<TValue> = [TValue | null, (newValue: LocalStorageSetStateValue<TValue> | null) => void, () => void];
+export type LocalStorageReturnValue<TValue> = [TValue, (newValue: LocalStorageSetStateValue<TValue> | null) => void, () => void];
 
 /**
  * React hook to enable updates to state via localStorage.
@@ -103,7 +104,7 @@ export function useLocalStorage<TValue = string>(
     };
   }, [key, defaultValue, onLocalStorageChange]);
 
-  const writeState = useCallback((value: TValue) => writeStorage(key, value), [key]);
+  const writeState = useCallback((value: LocalStorageSetStateValue<TValue>) => value instanceof Function ? writeStorage(key, value(localState)) : writeStorage(key, value), [key]);
   const deleteState = useCallback(() => deleteFromStorage(key), [key]);
   const state: TValue | null = localState ?? defaultValue;
 

test/use-localstorage.test.ts

@@ -2,11 +2,11 @@ import { useLocalStorage, deleteFromStorage } from '../src';
 import { renderHook, act } from '@testing-library/react-hooks';
 
 jest.mock('../src/storage', () => ({
-  ...jest.requireActual('../src/storage'),
-  storage: jest.fn(),
+    ...jest.requireActual('../src/storage'),
+    storage: jest.fn(),
 }));
 
-import { storage, LocalStorageProxy, MemoryStorageProxy  } from '../src/storage'
+import { storage, LocalStorageProxy, MemoryStorageProxy } from '../src/storage'
 
 describe('Module: use-localstorage', () => {
 
@@ -130,7 +130,7 @@ describe('Module: use-localstorage', () => {
                         const defaultValue = 'i';
                         const newValue = 'o';
                         const { result } = renderHook(() => useLocalStorage(key, defaultValue));
-                        
+
                         expect(result.current[0]).toBe(defaultValue);
                         expect(localStorage.getItem(key)).toBe(defaultValue);
 
@@ -144,6 +144,31 @@ describe('Module: use-localstorage', () => {
                     });
                 });
             });
+
+            describe('when a functional update is called', () => {
+                it('passes the current state to the update function', async () => {
+                    const key = 'functionUpdate';
+                    const defaultValue = 'ImTheDefault';
+                    const newValue = 'ImTheNewValue';
+                    const { result } = renderHook(() => useLocalStorage(key, defaultValue));
+
+                    act(() => result.current[1](prevValue => {
+                        expect(prevValue).toBe(defaultValue);
+                        return newValue;
+                    }));
+                    expect(result.current[0]).toBe(newValue);
+                });
+
+                it('handles updates to arrays', async () => {
+                    const key = 'arrayUpdate';
+                    const defaultValue = ['A'];
+                    const newValue = 'B';
+                    const { result } = renderHook(() => useLocalStorage(key, defaultValue));
+
+                    act(() => result.current[1](prevValue => prevValue ? [...prevValue, newValue] : [newValue]));
+                    expect(result.current[0]).toEqual([...defaultValue, newValue]);
+                })
+            });
         })
 
         describe("when localStorage api is disabled", () => {