Fix union as enum issues (#328)

.chronus/changes/fix_union_variant-2024-1-28-18-32-18.md

@@ -0,0 +1,7 @@
+---
+changeKind: fix
+packages:
+  - "@azure-tools/typespec-client-generator-core"
+---
+
+add support for anonymous union as enum and fix union as enum variant discriminator typing problem

packages/typespec-client-generator-core/src/interfaces.ts

@@ -161,6 +161,7 @@ export interface SdkDictionaryType extends SdkTypeBase {
 export interface SdkEnumType extends SdkTypeBase {
   kind: "enum";
   name: string;
+  generatedName?: string;
   valueType: SdkBuiltInType;
   values: SdkEnumValueType[];
   isFixed: boolean;

packages/typespec-client-generator-core/src/types.ts

@@ -450,11 +450,7 @@ function addDiscriminatorToModelType(
       if (discriminatorProperty.type.kind === "constant") {
         discriminatorType = { ...discriminatorProperty.type.valueType };
       } else if (discriminatorProperty.type.kind === "enumvalue") {
-        discriminatorType = getSdkEnum(
-          context,
-          (discriminatorProperty.type.__raw as EnumMember).enum,
-          operation
-        );
+        discriminatorType = discriminatorProperty.type.enumType;
       }
     } else {
       discriminatorType = {
@@ -650,11 +646,12 @@ function getSdkUnionEnum(context: TCGCContext, type: UnionEnum, operation?: Oper
     sdkType = {
       ...getSdkTypeBaseHelper(context, type.union, "enum"),
       name: getLibraryName(context, type.union),
+      generatedName: type.union.name ? undefined : getGeneratedName(context, type.union),
       description: docWrapper.description,
       details: docWrapper.details,
-      valueType: { ...getSdkTypeBaseHelper(context, type.kind, "string"), encode: "string" },
+      valueType: getSdkEnumValueType(context, type.flattenedMembers.values().next().value),
       values: [],
-      nullable: false,
+      nullable: type.nullable,
       isFixed: !type.open,
       isFlags: false,
       usage: UsageFlags.None, // We will add usage as we loop through the operations
@@ -761,9 +758,8 @@ export function getClientTypeWithDiagnostics(
       retval = getSdkEnum(context, type, operation);
       break;
     case "Union":
-      // start off with just handling nullable type
       const unionAsEnum = diagnostics.pipe(getUnionAsEnum(type));
-      if (unionAsEnum && type.name) {
+      if (unionAsEnum) {
         retval = getSdkUnionEnum(context, unionAsEnum, operation);
       } else {
         retval = diagnostics.pipe(getSdkUnionWithDiagnostics(context, type, operation));
@@ -778,7 +774,14 @@ export function getClientTypeWithDiagnostics(
       retval = getKnownValuesEnum(context, type, operation) ?? innerType;
       break;
     case "UnionVariant":
-      retval = diagnostics.pipe(getClientTypeWithDiagnostics(context, type.type, operation));
+      const unionType = diagnostics.pipe(
+        getClientTypeWithDiagnostics(context, type.union, operation)
+      );
+      if (unionType.kind === "enum") {
+        retval = unionType.values.find((x) => x.name === getLibraryName(context, type))!;
+      } else {
+        retval = diagnostics.pipe(getClientTypeWithDiagnostics(context, type.type, operation));
+      }
       break;
     case "EnumMember":
       const enumType = getSdkEnum(context, type.enum, operation);

packages/typespec-client-generator-core/test/public-utils.test.ts

@@ -1139,25 +1139,21 @@ describe("typespec-client-generator-core: public-utils", () => {
         `
         );
         const models = getAllModels(runner.context);
-        strictEqual(models.length, 1);
+        strictEqual(models.length, 2);
         const unionName = ((models[0] as SdkModelType).properties[0].type as SdkUnionType)
           .generatedName;
         strictEqual(unionName, "AStatus");
         strictEqual(models[0].kind, "model");
         const statusProp = models[0].properties[0];
         strictEqual(statusProp.kind, "property");
-        strictEqual(statusProp.type.kind, "union");
+        strictEqual(statusProp.type.kind, "enum");
         strictEqual(statusProp.type.values.length, 2);
-        const startVal = statusProp.type.values.find(
-          (x) => x.kind === "constant" && x.value === "start"
-        )!;
-        strictEqual(startVal.kind, "constant");
+        const startVal = statusProp.type.values.find((x) => x.name === "start")!;
+        strictEqual(startVal.kind, "enumvalue");
         strictEqual(startVal.valueType.kind, "string");
 
-        const stopVal = statusProp.type.values.find(
-          (x) => x.kind === "constant" && x.value === "stop"
-        )!;
-        strictEqual(stopVal.kind, "constant");
+        const stopVal = statusProp.type.values.find((x) => x.name === "stop")!;
+        strictEqual(stopVal.kind, "enumvalue");
         strictEqual(stopVal.valueType.kind, "string");
       });
 
@@ -1196,7 +1192,7 @@ describe("typespec-client-generator-core: public-utils", () => {
         `
         );
         const models = getAllModels(runner.context);
-        strictEqual(models.length, 2);
+        strictEqual(models.length, 3);
         const test1 = models.find((x) => (x as SdkModelType).generatedName === "AChoice")!;
         ok(test1);
         const unionName = ((test1 as SdkModelType).properties[0].type as SdkUnionType)

packages/typespec-client-generator-core/test/types.test.ts

@@ -1427,7 +1427,9 @@ describe("typespec-client-generator-core: types", () => {
     it("union to extensible enum values", async () => {
       await runner.compileWithBuiltInService(`
       union PetKind {
+        @doc("Cat")
         Cat: "cat",
+        @doc("Dog")
         Dog: "dog",
         string,
       }
@@ -1450,17 +1452,40 @@ describe("typespec-client-generator-core: types", () => {
 
       const catValue = values.find((x) => x.name === "Cat")!;
       strictEqual(catValue.value, "cat");
+      strictEqual(catValue.description, "Cat");
       strictEqual(catValue.enumType, petKind);
       strictEqual(catValue.valueType, petKind.valueType);
       strictEqual(catValue.kind, "enumvalue");
 
       const dogValue = values.find((x) => x.name === "Dog")!;
       strictEqual(dogValue.value, "dog");
+      strictEqual(dogValue.description, "Dog");
       strictEqual(dogValue.enumType, petKind);
       strictEqual(dogValue.valueType, petKind.valueType);
       strictEqual(dogValue.kind, "enumvalue");
     });
 
+    it("property of anonymous union as enum", async () => {
+      await runner.compileWithBuiltInService(`
+      model Pet {
+        kind: string | "cat" | "dog";
+      }
+
+      @route("/extensible-enum")
+      @put
+      op putPet(@body pet: Pet): void;
+      `);
+      const models = getAllModels(runner.context);
+      strictEqual(models.length, 2);
+      const pet = models.find((x) => x.name === "Pet")! as SdkModelType;
+      const kind = models.find((x) => x.name === "")!;
+      strictEqual(kind.generatedName, "PetKind");
+      const kindProperty = pet.properties.find(
+        (x) => (x.nameInClient = "kind")
+      )! as SdkBodyModelPropertyType;
+      strictEqual(kindProperty.type, kind);
+    });
+
     it("enum discriminator model without base discriminator property", async () => {
       await runner.compileWithBuiltInService(`
       enum DogKind {
@@ -1568,11 +1593,20 @@ describe("typespec-client-generator-core: types", () => {
       op getModel(): Fish;
       `);
       const models = getAllModels(runner.context);
-      strictEqual(models.length, 3);
+      strictEqual(models.length, 4);
+      const fish = models.find((x) => x.name === "Fish")! as SdkModelType;
+      let kindTypeProperty = fish.properties.find((x) => x.nameInClient === "kind")!;
+      strictEqual(kindTypeProperty.type.kind, "enum");
       const shark = models.find((x) => x.name === "Shark")! as SdkModelType;
       strictEqual(shark.discriminatorValue, "shark");
+      kindTypeProperty = shark.properties.find((x) => x.nameInClient === "kind")!;
+      strictEqual(kindTypeProperty.type.kind, "enumvalue");
       const salmon = models.find((x) => x.name === "Salmon")! as SdkModelType;
+      kindTypeProperty = salmon.properties.find((x) => x.nameInClient === "kind")!;
+      strictEqual(kindTypeProperty.type.kind, "enumvalue");
       strictEqual(salmon.discriminatorValue, "salmon");
+      const kindType = models.find((x) => x.name === "KindType")! as SdkEnumType;
+      strictEqual(kindType.isFixed, false);
     });
 
     it("filterOutCoreModels true", async () => {