example.ts

import { createStore } from "tinybase";
import z from "zod";
import { createTypedStore } from "./typed";

const exampleSchema = z.object({
  string: z.string(),
  nullableString: z.string().nullable(),
  number: z.number(),
  enum: z.enum(["a", "b"]),
  object: z.object({
    a: z.string(),
  }),
  array: z.array(z.string()),
});

const schema = {
  tables: {
    examples: exampleSchema,
  },
} as const;

const store = createStore();
const typedStore = createTypedStore(store, schema);

typedStore.setRow("examples", "1", {
  string: "s",
  nullableString: "s",
  number: 1,
  enum: "a",
  object: { a: "s" },
  array: ["a"],
});

const encodedRow = store.getRow("examples", "1");
/* complex cells are json encoded
{
  string: "s",
  nullableString: "s",
  number: 1,
  enum: "a",
  object: "{\"a\":\"s\"}",
  array: "[\"a\"]",
}
*/

const decodedRow = typedStore.getRow("examples", "1");
/* decoded to complex types
{
  string: "s",
  nullableString: "s",
  number: 1,
  enum: "a",
  object: {
    a: "s",
  },
  array: [ "a" ],
}
*/

// works for cells too
typedStore.setCell("examples", "1", "object", { a: "b" });
const cell = typedStore.getCell("examples", "1", "object");
// { a: "b" }

typed.ts

import { type Store } from "tinybase";
import * as z from "zod";
import type { $ZodShape, SomeType } from "zod/v4/core";

type StorageType = z.ZodString | z.ZodNumber | z.ZodBoolean | z.ZodNull;

type MapCellSchema<T extends SomeType> = T extends StorageType
  ? T
  : T extends z.ZodNullable<infer U>
  ? U extends StorageType
    ? T
    : z.ZodCodec<z.ZodString, T>
  : z.ZodCodec<z.ZodString, T>;

type AnyTablesSchema = Record<string, z.ZodObject>;

type AnySchema = {
  tables: AnyTablesSchema;
};

type MapRowShape<T extends $ZodShape> = {
  [K in keyof T]: MapCellSchema<T[K]>;
};

type MapRowSchema<T extends z.ZodObject> = T extends z.ZodObject<infer U>
  ? z.ZodObject<MapRowShape<U>>
  : never;

type MapTablesSchema<T extends AnyTablesSchema> = {
  [K in keyof T]: MapRowSchema<T[K]>;
};

type MapSchema<T extends AnySchema> = {
  tables: MapTablesSchema<T["tables"]>;
};

function mapJson(schema: SomeType) {
  return z.codec(z.string(), schema, {
    decode: (json) => JSON.parse(json),
    encode: (obj) => JSON.stringify(obj),
  });
}

function mapCell(schema: SomeType) {
  const type = schema._zod.def.type;
  if (type === "array" || type === "object") {
    return mapJson(schema) as any;
  }
  return schema;
}

function mapRow(schema: z.ZodObject) {
  return z.object(
    Object.fromEntries(
      Object.entries(schema.def.shape).map(([key, value]) => [
        key,
        mapCell(value),
      ])
    )
  );
}

function mapTables(schema: AnyTablesSchema) {
  return Object.fromEntries(
    Object.entries(schema).map(([tableId, rowSchema]) => [
      tableId,
      mapRow(rowSchema),
    ])
  );
}

function mapSchema(schema: AnySchema) {
  const tables = mapTables(schema.tables);
  return { tables };
}

export function createTypedStore<Schema extends AnySchema>(
  store: Store,
  schema: Schema
) {
  const mappedSchema = mapSchema(schema) as any;

  type MappedSchema = MapSchema<Schema>;
  type AnyTableId = Extract<keyof MappedSchema["tables"], string>;
  type RowType<TableId extends AnyTableId> = z.infer<
    MappedSchema["tables"][TableId]
  >;
  type AnyCellId<TableId extends AnyTableId> = Extract<
    keyof RowType<TableId>,
    string
  >;
  type CellType<
    TableId extends AnyTableId,
    ColumnId extends AnyCellId<TableId>
  > = RowType<TableId>[ColumnId];

  function getRowSchema<TableId extends AnyTableId>(tableId: TableId) {
    return mappedSchema.tables[tableId];
  }

  function getCellSchema(tableId: string, cellId: string) {
    return mappedSchema.tables[tableId].shape[cellId];
  }

  function getRow<TableId extends AnyTableId>(
    tableId: TableId,
    rowId: string
  ): RowType<TableId> {
    const schema = getRowSchema(tableId);
    const row = store.getRow(tableId, rowId);
    return schema.decode(row);
  }

  function setRow<TableId extends AnyTableId>(
    tableId: TableId,
    rowId: string,
    row: RowType<TableId>
  ): Store {
    const schema = getRowSchema(tableId);
    const encoded = schema.encode(row);
    return store.setRow(tableId, rowId, encoded);
  }

  function getCell<
    TableId extends AnyTableId,
    CellId extends AnyCellId<TableId>
  >(
    tableId: TableId,
    rowId: string,
    cellId: CellId
  ): CellType<TableId, CellId> | undefined {
    const schema = getCellSchema(tableId, cellId);
    const cell = store.getCell(tableId, rowId, cellId);
    return schema.decode(cell);
  }

  function setCell<
    TableId extends AnyTableId,
    CellId extends AnyCellId<TableId>,
    Cell extends CellType<TableId, CellId>,
    MapCell extends (cell: Cell | undefined) => Cell
  >(
    tableId: TableId,
    rowId: string,
    cellId: CellId,
    cell: Cell | MapCell
  ): Store {
    const schema = getCellSchema(tableId, cellId);
    const encoded = schema.encode(cell);
    return store.setCell(tableId, rowId, cellId, encoded);
  }

  return {
    getRow,
    setRow,
    getCell,
    setCell,
  };
}