Fix handling decorators for union variants in JSONSchema emitter

microsoft · May 18, 2024 · ddf6616 · ddf6616
1 parent 46d50af
commit ddf6616
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 4 deletions.
diff --git a/packages/compiler/src/emitter-framework/type-emitter.ts b/packages/compiler/src/emitter-framework/type-emitter.ts
@@ -118,7 +118,7 @@ export type EmitterOutput<T> = EmitEntity<T> | Placeholder<T> | T;
  *   literals, and any type referenced inside anywhere inside the model and any
  *   of the referenced types' references.
  *
- * In both cases, context is an object. It strongly recommended that the context
+ * In both cases, context is an object. It's strongly recommended that the context
  * object either contain only primitive types, or else only reference immutable
  * objects.
  *

diff --git a/packages/json-schema/src/json-schema-emitter.ts b/packages/json-schema/src/json-schema-emitter.ts
@@ -1,7 +1,6 @@
 import {
   BooleanLiteral,
   DiagnosticTarget,
-  DuplicateTracker,
   Enum,
   EnumMember,
   IntrinsicType,
@@ -52,6 +51,7 @@ import {
   SourceFileScope,
   TypeEmitter,
 } from "@typespec/compiler/emitter-framework";
+import { DuplicateTracker } from "@typespec/compiler/utils";
 import { stringify } from "yaml";
 import {
   JsonSchemaDeclaration,
@@ -328,7 +328,14 @@ export class JsonSchemaEmitter extends TypeEmitter<Record<string, any>, JSONSche
   }
 
   unionVariant(variant: UnionVariant): EmitterOutput<object> {
-    return this.emitter.emitTypeReference(variant.type);
+    const variantType = this.emitter.emitTypeReference(variant.type);
+    compilerAssert(variantType.kind === "code", "Unexpected non-code result from emit reference");
+
+    const result = new ObjectBuilder(variantType.value);
+
+    this.#applyConstraints(variant, result);
+
+    return result;
   }
 
   modelPropertyReference(property: ModelProperty): EmitterOutput<object> {
@@ -513,7 +520,7 @@ export class JsonSchemaEmitter extends TypeEmitter<Record<string, any>, JSONSche
   }
 
   #applyConstraints(
-    type: Scalar | Model | ModelProperty | Union | Enum,
+    type: Scalar | Model | ModelProperty | Union | UnionVariant | Enum,
     schema: ObjectBuilder<unknown>
   ) {
     const applyConstraint = (fn: (p: Program, t: Type) => any, key: string) => {

diff --git a/packages/json-schema/test/unions.test.ts b/packages/json-schema/test/unions.test.ts
@@ -87,4 +87,40 @@ describe("emitting unions", () => {
     const Foo = schemas["Foo.json"];
     assert.strictEqual(Foo["x-foo"], true);
   });
+
+  it("handles decorators on variants", async () => {
+    const schemas = await emitSchema(`
+      @jsonSchema
+      union Foo {
+        @doc("doc text")
+        @summary("summary text")
+        @extension("x-key", Json<"x-value">)
+        bar: string;
+
+        @doc("other model doc")
+        @summary("other model summary")
+        @extension("x-key-2", Json<"x-value-2">)
+        baz: OtherModel;
+      }
+
+      model OtherModel {}
+    `);
+
+    const Foo = schemas["Foo.json"];
+
+    assert.deepStrictEqual(Foo.anyOf, [
+      {
+        description: "doc text",
+        title: "summary text",
+        type: "string",
+        "x-key": "x-value",
+      },
+      {
+        $ref: "OtherModel.json",
+        description: "other model doc",
+        title: "other model summary",
+        "x-key-2": "x-value-2",
+      },
+    ]);
+  });
 });