fix(#19, #20): support complex unions and literal unions (#25)

Author: mvhenten

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<string, unknown>)} 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, {});
 

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"],
+	});
 };
 
 /**

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;

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[];
 }