TypedLocalStorage.ts - Adds typing and validation with zod for your local storage access, discarding corrupted local storage entries and falling back when local storage isn't available

TypedLocalStorage.ts

import { useCallback, useEffect, useRef, useState } from "react";
import { z } from "zod";
import { getUniqueId } from "./getUniqueId";

// Some examples
const localStorageValidators = {
	auth: z.object({
		token: z.string(),
		expiresAt: z.iso.datetime(),
	}),
	currentOrganizationId: z.string().optional(),
	sidebar: z.object({ open: z.boolean() }),
} as const;

export type LocalStorageKeys = keyof typeof localStorageValidators;
export type LocalStorageValue<K extends LocalStorageKeys> = z.output<
	(typeof localStorageValidators)[K]
>;

const localStorageCache: {
	[K in LocalStorageKeys]?: LocalStorageValue<K> | null;
} = {};
const localStorageSubscribers: {
	[K in LocalStorageKeys]?: Array<{
		id: number;
		fn: (value: LocalStorageValue<K> | null) => void;
	}>;
} = {};
const localStorageMetaValidator = z.object({
	value: z.unknown(),
	expiresAt: z.string().optional().nullable(),
});

function removeItem<K extends LocalStorageKeys>(key: K) {
	localStorageCache[key] = null;
	localStorage.removeItem(key);

	const subs = localStorageSubscribers[key];
	if (subs) {
		for (const sub of subs) {
			sub.fn(null);
		}
	}
}

function getItemBase<K extends LocalStorageKeys>(
	key: K,
): LocalStorageValue<K> | null {
	const itemString = localStorage.getItem(key);
	if (!itemString) {
		return null;
	}
	let json = null;
	try {
		json = JSON.parse(itemString);
	} catch (err) {
		console.error(err);
		return null;
	}

	const metaSafeParseOutput = localStorageMetaValidator.safeParse(json);
	if (!metaSafeParseOutput.success) {
		console.error(
			`Error parsing typed local storage meta '${key}'`,
			json,
			metaSafeParseOutput.error,
		);
		return null;
	}
	const { expiresAt: expiresAtStr } = metaSafeParseOutput.data;
	const expiresAt = expiresAtStr ? new Date(expiresAtStr) : null;

	if (expiresAt && new Date() > expiresAt) {
		removeItem(key);
		return null;
	}

	const untypedValue = metaSafeParseOutput.data.value;
	const validator = localStorageValidators[key];
	const safeParseOutput = validator.safeParse(untypedValue);
	if (!safeParseOutput.success) {
		console.error(
			`Error parsing typed local storage '${key}'`,
			untypedValue,
			metaSafeParseOutput.error,
		);
		return null;
	}

	return safeParseOutput.data as LocalStorageValue<K>;
}

function getItem<K extends LocalStorageKeys>(
	key: K,
): LocalStorageValue<K> | null {
	if (localStorageCache[key]) {
		return localStorageCache[key];
	}
	const item = getItemBase(key);
	// @ts-expect-error TS bad
	localStorageCache[key] = item;
	return item;
}

function setItem<K extends LocalStorageKeys>(
	key: K,
	value: LocalStorageValue<K>,
	{ expiresAt }: { expiresAt?: Date } = {},
) {
	// @ts-expect-error TS bad
	localStorageCache[key] = value;
	localStorage.setItem(
		key,
		JSON.stringify({
			value,
			expiresAt: expiresAt ? expiresAt.toISOString() : null,
		}),
	);
	const subs = localStorageSubscribers[key];
	if (subs) {
		for (const sub of subs) {
			sub.fn(value);
		}
	}
}

function subscribe<K extends LocalStorageKeys>(
	key: K,
	subscriberFn: (value: LocalStorageValue<K> | null) => void,
) {
	const id = getUniqueId();
	if (!localStorageSubscribers[key]) {
		localStorageSubscribers[key] = [];
	}
	localStorageSubscribers[key].push({ id, fn: subscriberFn });

	return () => {
		if (!localStorageSubscribers[key]) {
			return;
		}
		const index = localStorageSubscribers[key].findIndex((s) => s.id === id);
		if (index === -1) {
			return;
		}
		localStorageSubscribers[key].splice(index, 1);
	};
}

export const TypedLocalStorage = {
	removeItem,
	getItem,
	setItem,
};

export const useTypedLocalStorage = <K extends LocalStorageKeys>(key: K) => {
	const currentValueRef = useRef<LocalStorageValue<K> | null>(null);
	const currentValueFetchedRef = useRef<boolean>(false);
	if (!currentValueFetchedRef.current) {
		currentValueRef.current = TypedLocalStorage.getItem(key);
		currentValueFetchedRef.current = true;
	}
	const [currentValue, setCurrentValueBase] = useState(currentValueRef.current);
	const setCurrentValue = useCallback((value: LocalStorageValue<K> | null) => {
		if (currentValueRef.current !== value) {
			currentValueRef.current = value;
			setCurrentValueBase(value);
		}
	}, []);
	const setValue = useCallback(
		(value: LocalStorageValue<K>, { expiresAt }: { expiresAt?: Date } = {}) => {
			setCurrentValue(value);
			TypedLocalStorage.setItem(key, value, { expiresAt });
		},
		[key, setCurrentValue],
	);
	const remove = useCallback(() => {
		setCurrentValue(null);
		TypedLocalStorage.removeItem(key);
	}, [key, setCurrentValue]);
	useEffect(() => {
		const unsub = subscribe(key, (value) => {
			setCurrentValue(value);
		});
		return unsub;
	}, [key, setCurrentValue]);
	return [currentValue, setValue, remove] as const;
};