Union enum (#156)

common/changes/@azure-tools/typespec-client-generator-core/unionEnum_2024-01-23-23-51.json

@@ -0,0 +1,10 @@
+{
+  "changes": [
+    {
+      "packageName": "@azure-tools/typespec-client-generator-core",
+      "comment": "add support for creating an enum from a union definition",
+      "type": "none"
+    }
+  ],
+  "packageName": "@azure-tools/typespec-client-generator-core"
+}

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

@@ -1,6 +1,4 @@
 import {
-  Enum,
-  EnumMember,
   getDeprecationDetails,
   getDoc,
   getEffectiveModelType,
@@ -9,13 +7,13 @@ import {
   getProjectedName,
   getSummary,
   ignoreDiagnostics,
+  Interface,
   listServices,
   Model,
   ModelProperty,
   Namespace,
   Operation,
   resolveEncodedName,
-  Scalar,
   Type,
   Union,
 } from "@typespec/compiler";
@@ -178,10 +176,7 @@ export function getPropertyNames(context: SdkContext, property: ModelProperty):
  * @param type
  * @returns the library name for a typespec type
  */
-export function getLibraryName(
-  context: SdkContext,
-  type: Model | ModelProperty | Operation | Enum | EnumMember
-): string {
+export function getLibraryName(context: SdkContext, type: Type & { name?: string }): string {
   // 1. check if there's a client name
   let emitterSpecificName = getClientNameOverride(context, type);
   if (emitterSpecificName) return emitterSpecificName;
@@ -235,8 +230,8 @@ export function getWireName(context: SdkContext, type: Type & { name: string })
   return getProjectedName(context.program, type, "json") ?? type.name;
 }
 
-interface DefaultSdkTypeBase<TKind> {
-  __raw: Type;
+interface SdkTypeBaseHelper<TKind> {
+  __raw?: Type;
   nullable: boolean;
   deprecation?: string;
   kind: TKind;
@@ -248,9 +243,15 @@ interface DefaultSdkTypeBase<TKind> {
  */
 export function getSdkTypeBaseHelper<TKind>(
   context: SdkContext,
-  type: Type,
+  type: Type | string,
   kind: TKind
-): DefaultSdkTypeBase<TKind> {
+): SdkTypeBaseHelper<TKind> {
+  if (typeof type === "string") {
+    return {
+      nullable: false,
+      kind,
+    };
+  }
   return {
     __raw: type,
     nullable: false,
@@ -264,7 +265,12 @@ export function getSdkTypeBaseHelper<TKind>(
  * @param type
  * @returns
  */
-export function getCrossLanguageDefinitionId(type: Model | Enum | Operation | Scalar): string {
+export function getCrossLanguageDefinitionId(type: {
+  name: string;
+  kind: string;
+  interface?: Interface;
+  namespace?: Namespace;
+}): string {
   let retval = type.name;
   if (type.kind === "Operation" && type.interface) {
     retval = `${type.interface.name}.${retval}`;

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

@@ -1,4 +1,9 @@
-import { getLroMetadata, isFixed } from "@azure-tools/typespec-azure-core";
+import {
+  UnionEnum,
+  getLroMetadata,
+  getUnionAsEnum,
+  isFixed,
+} from "@azure-tools/typespec-azure-core";
 import {
   BooleanLiteral,
   BytesKnownEncoding,
@@ -531,6 +536,53 @@ export function getSdkEnum(context: SdkContext, type: Enum, operation?: Operatio
   return sdkType;
 }
 
+function getSdkUnionEnumValues(
+  context: SdkContext,
+  type: UnionEnum,
+  enumType: SdkEnumType
+): SdkEnumValueType[] {
+  const values: SdkEnumValueType[] = [];
+  for (const [name, member] of type.flattenedMembers.entries()) {
+    const docWrapper = getDocHelper(context, member.variant);
+    values.push({
+      kind: "enumvalue",
+      name: typeof name === "string" ? name : `${member.value}`,
+      description: docWrapper.description,
+      details: docWrapper.details,
+      value: member.value,
+      valueType: enumType.valueType,
+      enumType,
+      nullable: false,
+    });
+  }
+  return values;
+}
+
+function getSdkUnionEnum(context: SdkContext, type: UnionEnum, operation?: Operation) {
+  let sdkType = context.modelsMap?.get(type.union) as SdkEnumType | undefined;
+  if (!sdkType) {
+    const union = type.union as Union & { name: string };
+    const docWrapper = getDocHelper(context, union);
+    sdkType = {
+      ...getSdkTypeBaseHelper(context, type.union, "enum"),
+      name: getLibraryName(context, type.union),
+      description: docWrapper.description,
+      details: docWrapper.details,
+      valueType: { ...getSdkTypeBaseHelper(context, type.kind, "string"), encode: "string" },
+      values: [],
+      nullable: false,
+      isFixed: !type.open,
+      isFlags: false,
+      usage: UsageFlags.None, // We will add usage as we loop through the operations
+      access: undefined, // Dummy value until we update models map
+      crossLanguageDefinitionId: getCrossLanguageDefinitionId(union),
+    };
+    sdkType.values = getSdkUnionEnumValues(context, type, sdkType);
+  }
+  updateModelsMap(context, type.union, sdkType, operation);
+  return sdkType;
+}
+
 function getKnownValuesEnum(
   context: SdkContext,
   type: Scalar | ModelProperty,
@@ -606,6 +658,10 @@ export function getClientType(context: SdkContext, type: Type, operation?: Opera
       return getSdkEnum(context, type, operation);
     case "Union":
       // start off with just handling nullable type
+      const unionAsEnum = ignoreDiagnostics(getUnionAsEnum(type));
+      if (unionAsEnum && type.name) {
+        return getSdkUnionEnum(context, unionAsEnum, operation);
+      }
       const union = getSdkUnion(context, type, operation);
       if (union === undefined) {
         throw Error(`Error encountered during generation, view diagnostic logs`);

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

@@ -1370,6 +1370,43 @@ describe("typespec-client-generator-core: types", () => {
       strictEqual(dogKindProperty.type, dogKind);
     });
 
+    it("union to extensible enum values", async () => {
+      await runner.compileWithBuiltInService(`
+      union PetKind {
+        Cat: "cat",
+        Dog: "dog",
+        string,
+      }
+
+      @route("/extensible-enum")
+      @put
+      op putPet(@body petKind: PetKind): void;
+      `);
+      const models = Array.from(getAllModels(runner.context));
+      strictEqual(models.length, 1);
+      const petKind = models[0] as SdkEnumType;
+      strictEqual(petKind.name, "PetKind");
+      strictEqual(petKind.isFixed, false);
+      strictEqual(petKind.valueType.kind, "string");
+      const values = petKind.values;
+      deepStrictEqual(
+        values.map((x) => x.name),
+        ["Cat", "Dog"]
+      );
+
+      const catValue = values.find((x) => x.name === "Cat")!;
+      strictEqual(catValue.value, "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.enumType, petKind);
+      strictEqual(dogValue.valueType, petKind.valueType);
+      strictEqual(dogValue.kind, "enumvalue");
+    });
+
     it("enum discriminator model without base discriminator property", async () => {
       await runner.compileWithBuiltInService(`
       enum DogKind {