Pinia's define store with QOL improvements

defineStore.ts

import type { StateTree, Store } from 'pinia';
import { defineStore as defineStoreBase } from 'pinia';
import type { ComputedRef } from 'vue';
import { isRef, toRaw, toRef } from 'vue';

// Creates a store without wrapping the returned values in reactive()
// to allow easy destructuring in components
// Removes extra methods from it, so stores can be easily combined
// The original store is kept in store._inner
export function defineStore<Id extends string, SS extends StateTree>(
    name: Id,
    setupStore: () => SS
): StoreDefinitionUnwrapped<Id, SS> {
    const storeDefinition = defineStoreBase(name, setupStore);
    const unwrappedStoreDefinition = () => cleanupStore(unwrapStore(storeDefinition()));

    return setInnerStore(unwrappedStoreDefinition, storeDefinition);
}

function unwrapStore<T extends Record<string, unknown>>(store: T): T {
    const rawStore = toRaw(store);

    const unwrappedStore = objFromEntries(
        objEntries(rawStore).map(([key, value]) => {
            return [key, isRef(value) ? toRef(rawStore, key) : rawStore[key]];
        })
    );

    return unwrappedStore as T;
}

const extraKeys: Record<keyof Store, boolean> = {
    $state: true,
    $patch: true,
    $reset: true,
    $subscribe: true,
    $onAction: true,
    $dispose: true,
    $id: true,
    _customProperties: true,
};

function cleanupStore<T extends Record<string, unknown>>(store: T): Omit<T, keyof typeof extraKeys> {
    const storeCopy = { ...store };

    objKeys(extraKeys).forEach((key) => {
        delete storeCopy[key];
    });

    return storeCopy;
}

function setInnerStore<T extends StoreDefinitionUnwrapped<any, any>>(
    storeDefinition: Omit<T, '_inner'> & Partial<Pick<T, '_inner'>>,
    inner: T['_inner']
): T {
    storeDefinition._inner = inner;

    return storeDefinition as T;
}

type StoreDefinitionUnwrapped<Id extends string, SS extends StateTree> = {
    (): StoreUnwrapped<Id, ExtractState<SS>, ExtractGetters<SS>, ExtractActions<SS>>;
    _inner: ReturnType<typeof defineStoreBase<Id, SS>>;
};

// Removed a bunch of extra stuff here, it can be used via _inner if need be
type StoreUnwrapped<Id extends string, S extends StateTree, G, A> = S & G & A;

/* Updated internal pinia types */

// Removed UnwrapAll here
type ExtractState<SS> = SS extends undefined | void
    ? {}
    : ExtractStateKeys<SS> extends keyof SS
    ? Pick<SS, ExtractStateKeys<SS>>
    : never;

type ExtractStateKeys<SS> = keyof {
    [K in keyof SS as SS[K] extends ((...args: any[]) => any) | ComputedRef ? never : K]: any;
};

type ExtractActions<SS> = SS extends undefined | void
    ? {}
    : ExtractActionsKeys<SS> extends keyof SS
    ? Pick<SS, ExtractActionsKeys<SS>>
    : never;

type ExtractActionsKeys<SS> = keyof {
    [K in keyof SS as SS[K] extends (...args: any[]) => any ? K : never]: any;
};

type ExtractGetters<SS> = SS extends undefined | void
    ? {}
    : ExtractGettersKeys<SS> extends keyof SS
    ? Pick<SS, ExtractGettersKeys<SS>>
    : never;

type ExtractGettersKeys<SS> = keyof {
    [K in keyof SS as SS[K] extends ComputedRef ? K : never]: any;
};

/* Typed utils */

type Entries<T> = Array<[keyof T, Exclude<T[keyof T], undefined>]>;

function objEntries<T extends Record<PropertyKey, unknown>>(object: T): Entries<T> {
    return Object.entries(object) as Entries<T>;
}

function objKeys<T extends Record<PropertyKey, unknown>>(object: T): Array<keyof T> {
    return Object.keys(object);
}

function objFromEntries<T extends PropertyKey, U>(
    entries: ReadonlyArray<readonly [T, U]>
): Readonly<{ [K in T]: U }> {
    return Object.fromEntries(entries) as { [K in T]: U };
}