nicer unions

2 files changed, 10 insertions, 5 deletions

diff --git a/src/util/codec-test.ts b/src/util/codec-test.ts
index 0d1ab5603..640cd7c1c 100644
--- a/src/util/codec-test.ts
+++ b/src/util/codec-test.ts
@@ -64,7 +64,8 @@ test("union", t => {
     .property("type", stringConstCodec("two"))
     .property("bar", stringCodec)
     .build("AltTwo");
-  const myUnionCodec: Codec<MyUnion> = unionCodec<MyUnion, "type">("type")
+  const myUnionCodec: Codec<MyUnion> = unionCodec<MyUnion>()
+    .discriminateOn("type")
     .alternative("one", altOneCodec)
     .alternative("two", altTwoCodec)
     .build<MyUnion>("MyUnion");
diff --git a/src/util/codec.ts b/src/util/codec.ts
index 78516183c..a6fe74905 100644
--- a/src/util/codec.ts
+++ b/src/util/codec.ts
@@ -233,6 +233,12 @@ export function mapCodec<T>(innerCodec: Codec<T>): Codec<{ [x: string]: T }> {
   };
 }
 
+export class UnionCodecPreBuilder<T> {
+  discriminateOn<D extends keyof T>(discriminator: D): UnionCodecBuilder<T, D, never> {
+    return new UnionCodecBuilder<T, D, never>(discriminator);
+  }
+}
+
 /**
  * Return a builder for a codec that decodes an object with properties.
  */
@@ -240,8 +246,6 @@ export function objectCodec<T>(): ObjectCodecBuilder<T, {}> {
   return new ObjectCodecBuilder<T, {}>();
 }
 
-export function unionCodec<T, D extends keyof T>(
-  discriminator: D,
-): UnionCodecBuilder<T, D, never> {
-  return new UnionCodecBuilder<T, D, never>(discriminator);
+export function unionCodec<T>(): UnionCodecPreBuilder<T> {
+  return new UnionCodecPreBuilder<T>();
 }