Allow encode on a union (#2992)

fix [#2969](#2969)
microsoft · Mar 12, 2024 · 128a508 · 128a508
1 parent d8576ef
commit 128a508
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 13 deletions.
diff --git a/.chronus/changes/allow-encode-union-2024-2-6-17-42-29.md b/.chronus/changes/allow-encode-union-2024-2-6-17-42-29.md
@@ -0,0 +1,8 @@
+---
+# Change versionKind to one of: internal, fix, dependencies, feature, deprecation, breaking
+changeKind: feature
+packages:
+  - "@typespec/compiler"
+---
+
+Enable the use of `@encode` for model properties that have a union type. This supports cases like `@encode("rfc3339") prop: utcDateTime | null`
diff --git a/packages/compiler/src/lib/decorators.ts b/packages/compiler/src/lib/decorators.ts
@@ -756,25 +756,20 @@ export function $encode(
     type: encodeAs ?? context.program.checker.getStdType("string"),
   };
   const targetType = getPropertyType(target);
-  if (targetType.kind !== "Scalar") {
-    return;
-  }
   validateEncodeData(context, targetType, encodeData);
   context.program.stateMap(encodeKey).set(target, encodeData);
 }
 
-function validateEncodeData(context: DecoratorContext, target: Scalar, encodeData: EncodeData) {
+function validateEncodeData(context: DecoratorContext, target: Type, encodeData: EncodeData) {
   function check(validTargets: StdTypeName[], validEncodeTypes: StdTypeName[]) {
     const checker = context.program.checker;
-    const isTargetValid = validTargets.some((validTarget) => {
-      return ignoreDiagnostics(
-        checker.isTypeAssignableTo(
-          target.projectionBase ?? target,
-          checker.getStdType(validTarget),
-          target
-        )
-      );
-    });
+    const isTargetValid = isTypeIn(target.projectionBase ?? target, (type) =>
+      validTargets.some((validTarget) => {
+        return ignoreDiagnostics(
+          checker.isTypeAssignableTo(type, checker.getStdType(validTarget), target)
+        );
+      })
+    );
 
     if (!isTargetValid) {
       reportDiagnostic(context.program, {

diff --git a/packages/compiler/test/decorators/decorators.test.ts b/packages/compiler/test/decorators/decorators.test.ts
@@ -587,6 +587,30 @@ describe("compiler: built-in decorators", () => {
       strictEqual(getEncode(runner.program, s)?.encoding, "rfc3339");
     });
 
+    it(`set encoding on model property`, async () => {
+      const { prop } = (await runner.compile(`
+        model Foo {
+          @encode("rfc3339")
+          @test
+          prop: utcDateTime;
+        }
+      `)) as { prop: ModelProperty };
+
+      strictEqual(getEncode(runner.program, prop)?.encoding, "rfc3339");
+    });
+
+    it(`set encoding on model property of union type`, async () => {
+      const { prop } = (await runner.compile(`
+        model Foo {
+          @encode("rfc3339")
+          @test
+          prop: utcDateTime | null; 
+        }
+      `)) as { prop: ModelProperty };
+
+      strictEqual(getEncode(runner.program, prop)?.encoding, "rfc3339");
+    });
+
     it(`encode type default to string`, async () => {
       const { s } = (await runner.compile(`
         @encode("rfc3339")