diff --git a/src/emitter.ts b/src/emitter.ts index b3439e8..6897c6d 100644 --- a/src/emitter.ts +++ b/src/emitter.ts @@ -7,6 +7,7 @@ import type { RecordModelType, Scalar, Type, + Union, } from "@typespec/compiler"; import { type EmitContext, @@ -18,7 +19,7 @@ import type { Attribute, CustomAttribute, Schema } from "electrodb"; import * as ts from "typescript"; import { StateKeys } from "./lib.js"; -import { stringifyObject } from "./stringify.js"; +import { RawCode, stringifyObject } from "./stringify.js"; function emitIntrinsincScalar(type: Scalar) { switch (type.name) { @@ -109,6 +110,105 @@ function emitEnumModel(type: Enum): Attribute { }; } +function emitTypeToTypeScript(type: Type): string { + switch (type.kind) { + case "Scalar": { + let baseType = type; + while (baseType.baseScalar) { + baseType = baseType.baseScalar; + } + switch (baseType.name) { + case "boolean": + return "boolean"; + case "numeric": + case "integer": + case "float": + case "int64": + case "int32": + case "int16": + case "int8": + case "uint64": + case "uint32": + case "uint16": + case "uint8": + case "safeint": + case "float32": + case "float64": + case "decimal": + case "decimal128": + return "number"; + default: + return "string"; + } + } + case "Model": { + if (type.name === "Array") { + const arrayType = type as ArrayModelType; + return `${emitTypeToTypeScript(arrayType.indexer.value)}[]`; + } + const properties: string[] = []; + for (const prop of walkPropertiesInherited(type as RecordModelType)) { + const optional = prop.optional ? "?" : ""; + properties.push( + `${prop.name}${optional}: ${emitTypeToTypeScript(prop.type)}`, + ); + } + return `{ ${properties.join("; ")} }`; + } + case "Enum": { + const values = Array.from(type.members) + .map(([key, member]) => `"${member.value ?? key}"`) + .join(" | "); + return values; + } + case "Union": { + const variants = Array.from(type.variants.values()) + .map((variant) => emitTypeToTypeScript(variant.type)) + .join(" | "); + return variants; + } + default: + return "any"; + } +} + +function isLiteralUnion(type: Union): string[] | null { + const literals: string[] = []; + + for (const variant of type.variants.values()) { + // Check if this variant is a string or number literal + if (variant.type.kind === "String") { + literals.push(variant.type.value); + } else if (variant.type.kind === "Number") { + literals.push(String(variant.type.value)); + } else { + // Not a literal union, return null + return null; + } + } + + return literals; +} + +function emitUnion(type: Union): Attribute { + // Check if this is a simple literal union (e.g., "home" | "work" | "other") + const literals = isLiteralUnion(type); + if (literals) { + // Emit as enum-like array, similar to how named enums are handled + return { + type: literals, + }; + } + + // Complex union - use CustomAttributeType + const tsType = emitTypeToTypeScript(type); + // RawCode is used to emit the CustomAttributeType function call as-is + return { + // @ts-expect-error - RawCode is handled by stringifyObject at code generation time + type: new RawCode(`CustomAttributeType<${tsType}>("any")`), + }; +} + function emitType(type: Type): Attribute { switch (type.kind) { case "Scalar": @@ -118,7 +218,7 @@ function emitType(type: Type): Attribute { case "Enum": return emitEnumModel(type); case "Union": - return { type: "string" }; + return emitUnion(type); default: throw new Error(`Type kind ${type.kind} is currently not supported!`); } @@ -219,13 +319,23 @@ export async function $onEmit(context: EmitContext) { }; } - const typescriptSource = Object.entries(entities) + const entityDefinitions = Object.entries(entities) .map( ([name, schema]) => `export const ${name} = ${stringifyObject(schema as unknown as Record)} as const`, ) .join("\n"); + // Add CustomAttributeType import if any union types are used + const hasCustomAttributeType = entityDefinitions.includes( + "CustomAttributeType", + ); + const imports = hasCustomAttributeType + ? 'import { CustomAttributeType } from "electrodb";\n\n' + : ""; + + const typescriptSource = imports + entityDefinitions; + const declarations = await ts.transpileDeclaration(typescriptSource, {}); const javascript = await ts.transpileModule(typescriptSource, {}); diff --git a/src/stringify.ts b/src/stringify.ts index a3472c0..acfaff7 100644 --- a/src/stringify.ts +++ b/src/stringify.ts @@ -3,6 +3,13 @@ import { generate } from "@babel/generator"; import { parseExpression } from "@babel/parser"; import type { Expression, ObjectExpression } from "@babel/types"; +/** + * Marker class for raw code expressions that should be emitted as-is + */ +export class RawCode { + constructor(public readonly code: string) {} +} + function isObjectExpression(target: unknown): target is ObjectExpression { return (target as Expression | undefined)?.type === "ObjectExpression"; } @@ -26,6 +33,10 @@ const stringifyValue = (value: unknown) => { case "boolean": return value.toString(); case "object": + if (value instanceof RawCode) { + return value.code; + } + if (Array.isArray(value)) { return stringifyArray(value); } @@ -45,7 +56,9 @@ const stringifyArray = (value: unknown[]): string => { }; const stringifyKeyValue = (key: string, value: unknown) => { - return parseExpression(`{ ${key}: ${stringifyValue(value)} }`); + return parseExpression(`{ ${key}: ${stringifyValue(value)} }`, { + plugins: ["typescript"], + }); }; /** diff --git a/test/entities.test.js b/test/entities.test.js index 9c07dee..b50d555 100644 --- a/test/entities.test.js +++ b/test/entities.test.js @@ -169,9 +169,9 @@ suite("Person Entity", () => { }); }); - test("address.type property (union literal) is string", () => { + test("address.type property (literal union) has enum-like values", () => { assert.deepEqual(Person.attributes.address.properties.type, { - type: "string", + type: ["home", "work", "other"], required: true, }); }); @@ -228,6 +228,34 @@ suite("Person Entity", () => { }); }); + suite("Union type (Info[] with BooleanValue | Int64Value)", () => { + test("additionalInfo is a list type with required: true", () => { + assert.equal(Person.attributes.additionalInfo.type, "list"); + assert.equal(Person.attributes.additionalInfo.required, true); + }); + + test("additionalInfo items are map type", () => { + assert.equal(Person.attributes.additionalInfo.items.type, "map"); + }); + + test("additionalInfo item has name property as string", () => { + assert.deepEqual( + Person.attributes.additionalInfo.items.properties.name, + { + type: "string", + required: true, + }, + ); + }); + + test("additionalInfo item value property uses CustomAttributeType for union", () => { + const valueAttr = Person.attributes.additionalInfo.items.properties.value; + // CustomAttributeType("any") returns "any" at runtime + assert.equal(valueAttr.type, "any"); + assert.equal(valueAttr.required, true); + }); + }); + suite("Index configurations", () => { test("persons index (primary, pk only)", () => { const personsIndex = Person.indexes.persons; diff --git a/test/main.tsp b/test/main.tsp index 84b6a39..81a3c3b 100644 --- a/test/main.tsp +++ b/test/main.tsp @@ -24,6 +24,20 @@ enum CountryCode { US, DE, } +model Int64Value { + value: int64; + type: "int64"; +} + +model BooleanValue { + value: boolean; + type: "boolean"; +} + +model Info { + name: string; + value: BooleanValue | Int64Value; +} model Address { street: String64; @@ -106,5 +120,7 @@ model Person { status: PersonStatus; + additionalInfo: Info[]; + coffeePreferences: CoffeePreferences[]; }