handle certain kinds of union types more correctly

Author: mathpirate

packages/schema-generator/src/formatters/common-tools-formatter.ts

@@ -45,6 +45,10 @@ export class CommonToolsFormatter implements TypeFormatter {
       return true;
     }
 
+    if ((type.flags & ts.TypeFlags.Union) !== 0) {
+      return false;
+    }
+
     // Check if this is a wrapper type (Cell/Stream/OpaqueRef) via type structure
     const wrapperInfo = getCellWrapperInfo(type, context.typeChecker);
     return wrapperInfo !== undefined;
@@ -93,7 +97,7 @@ export class CommonToolsFormatter implements TypeFormatter {
     }
 
     const wrapperInfo = getCellWrapperInfo(type, context.typeChecker);
-    if (wrapperInfo) {
+    if (wrapperInfo && !(type.flags & ts.TypeFlags.Union)) {
       const nodeToPass = this.selectWrapperTypeNode(
         n,
         resolvedWrapper,
@@ -250,9 +254,11 @@ export class CommonToolsFormatter implements TypeFormatter {
     }
 
     // Check if this type itself is the target wrapper
-    const wrapperInfo = getCellWrapperInfo(type, checker);
-    if (wrapperInfo && wrapperInfo.kind === targetWrapperKind) {
-      return { type, typeRef: wrapperInfo.typeRef, kind: wrapperInfo.kind };
+    if ((type.flags & ts.TypeFlags.Union) === 0) {
+      const wrapperInfo = getCellWrapperInfo(type, checker);
+      if (wrapperInfo && wrapperInfo.kind === targetWrapperKind) {
+        return { type, typeRef: wrapperInfo.typeRef, kind: wrapperInfo.kind };
+      }
     }
 
     // If this is a union (e.g., from Opaque<T>), check each member

packages/schema-generator/test/fixtures/schema/cell-brand-optional.expected.json

@@ -0,0 +1,38 @@
+{
+  "$defs": {
+    "MyCell": {
+      "asCell": true,
+      "type": "number"
+    },
+    "NullableCell": {
+      "anyOf": [
+        {
+          "type": "null"
+        },
+        {
+          "$ref": "#/$defs/MyCell"
+        }
+      ]
+    },
+    "OptionalCell": {
+      "anyOf": [
+        {
+          "$ref": "#/$defs/MyCell"
+        }
+      ]
+    }
+  },
+  "properties": {
+    "maybeCell": {
+      "$ref": "#/$defs/OptionalCell"
+    },
+    "nullableCell": {
+      "$ref": "#/$defs/NullableCell"
+    }
+  },
+  "required": [
+    "maybeCell",
+    "nullableCell"
+  ],
+  "type": "object"
+}

packages/schema-generator/test/fixtures/schema/cell-brand-optional.input.ts

@@ -0,0 +1,8 @@
+interface MyCell<T> extends BrandedCell<T, "cell"> {}
+type OptionalCell = MyCell<number> | undefined;
+type NullableCell = MyCell<string> | null;
+
+interface SchemaRoot {
+  maybeCell: OptionalCell;
+  nullableCell: NullableCell;
+}

packages/utils/deno.json

@@ -1,7 +1,7 @@
 {
   "name": "@commontools/utils",
   "tasks": {
-    "test": "deno test"
+    "test": "deno test --allow-env"
   },
   "exports": {
     ".": "./src/index.ts",

packages/utils/src/typescript/cell-brand.ts

@@ -199,6 +199,14 @@ function extractWrapperTypeReference(
     }
   }
 
+  if (type.flags & ts.TypeFlags.Union) {
+    const unionType = type as ts.UnionType;
+    for (const member of unionType.types) {
+      const ref = extractWrapperTypeReference(member, checker, seen);
+      if (ref) return ref;
+    }
+  }
+
   return undefined;
 }
 

packages/utils/test/cell-brand.test.ts

@@ -0,0 +1,92 @@
+import { assert, assertEquals, assertExists } from "@std/assert";
+import ts from "typescript";
+
+import { getCellWrapperInfo } from "../src/typescript/cell-brand.ts";
+
+function createProgram(source: string): {
+  checker: ts.TypeChecker;
+  sourceFile: ts.SourceFile;
+} {
+  const fileName = "test.ts";
+  const sourceFile = ts.createSourceFile(
+    fileName,
+    source,
+    ts.ScriptTarget.ES2020,
+    true,
+    ts.ScriptKind.TS,
+  );
+
+  const compilerOptions: ts.CompilerOptions = {
+    target: ts.ScriptTarget.ES2020,
+    module: ts.ModuleKind.ESNext,
+    strict: true,
+  };
+
+  const host = ts.createCompilerHost(compilerOptions, true);
+
+  host.getSourceFile = (name) => name === fileName ? sourceFile : undefined;
+  host.readFile = (name) => name === fileName ? source : undefined;
+  host.fileExists = (name) => name === fileName;
+  host.getDirectories = () => [];
+  host.getCurrentDirectory = () => "/";
+  host.writeFile = () => {};
+
+  const program = ts.createProgram([fileName], compilerOptions, host);
+  return { checker: program.getTypeChecker(), sourceFile };
+}
+
+function getPropertyType(
+  checker: ts.TypeChecker,
+  sourceFile: ts.SourceFile,
+  interfaceName: string,
+  propertyName: string,
+): ts.Type {
+  const statements = sourceFile.statements.filter(ts.isInterfaceDeclaration);
+  const iface = statements.find((stmt) => stmt.name.text === interfaceName);
+  assertExists(iface, `Interface ${interfaceName} not found`);
+
+  const property = iface.members.find((member): member is ts.PropertySignature =>
+    ts.isPropertySignature(member) &&
+    !!member.name &&
+    ts.isIdentifier(member.name) &&
+    member.name.text === propertyName
+  );
+  assertExists(property, `Property ${propertyName} not found`);
+
+  return checker.getTypeAtLocation(property);
+}
+
+Deno.test("getCellWrapperInfo handles unions containing branded cells", () => {
+  const source = `
+    declare const CELL_BRAND: unique symbol;
+    interface BrandedCell<T, Brand extends string> {
+      readonly [CELL_BRAND]: Brand;
+    }
+    interface Cell<T> extends BrandedCell<T, "cell"> {}
+    type MaybeCell = Cell<number> | undefined;
+    type NullableCell = Cell<string> | null;
+
+    interface Schema {
+      maybe: MaybeCell;
+      nullable: NullableCell;
+    }
+  `;
+
+  const { checker, sourceFile } = createProgram(source);
+  const maybeType = getPropertyType(checker, sourceFile, "Schema", "maybe");
+  const nullableType = getPropertyType(checker, sourceFile, "Schema", "nullable");
+
+  const maybeInfo = getCellWrapperInfo(maybeType, checker);
+  assertExists(maybeInfo, "Expected wrapper info for MaybeCell");
+  assertEquals(maybeInfo.kind, "Cell");
+  const maybeArg = maybeInfo.typeRef.typeArguments?.[0];
+  assertExists(maybeArg);
+  assertEquals(checker.typeToString(maybeArg), "number");
+
+  const nullableInfo = getCellWrapperInfo(nullableType, checker);
+  assertExists(nullableInfo, "Expected wrapper info for NullableCell");
+  assertEquals(nullableInfo.kind, "Cell");
+  const nullableArg = nullableInfo.typeRef.typeArguments?.[0];
+  assertExists(nullableArg);
+  assertEquals(checker.typeToString(nullableArg), "string");
+});