Redesign discriminated union type parser (#2068)

Previously, the `discriminatedUnion` parser required you to specify the
discriminant via an object parser, which was confusing and verbose. Now you only
have to specify the discriminant key once, and provide the values for each
branch.

Issue: LEMS-2582

## Test plan:

`yarn test`

Author: benchristel

Reviewers: jeremywiebe, benchristel, anakaren-rojas, catandthemachines, nishasy

Required Reviewers:

Approved By: jeremywiebe

Checks: ✅ Publish npm snapshot (ubuntu-latest, 20.x), ✅ Cypress (ubuntu-latest, 20.x), ✅ Check for .changeset entries for all changed files (ubuntu-latest, 20.x), ✅ Check builds for changes in size (ubuntu-latest, 20.x), ✅ Lint, Typecheck, Format, and Test (ubuntu-latest, 20.x), ✅ Publish Storybook to Chromatic (ubuntu-latest, 20.x)

Pull Request URL: #2068

.changeset/many-cats-run.md

@@ -0,0 +1,5 @@
+---
+"@khanacademy/perseus": patch
+---
+
+Internal: Redesign discriminated union type parser to have a simpler and more intuitive interface.

packages/perseus/src/util/parse-perseus-json/general-purpose-parsers/discriminated-union.test.ts

@@ -2,83 +2,118 @@ import {parse} from "../parse";
 import {failure, success} from "../result";
 
 import {constant} from "./constant";
-import {discriminatedUnion} from "./discriminated-union";
+import {discriminatedUnionOn} from "./discriminated-union";
 import {number} from "./number";
 import {object} from "./object";
 
-describe("a discriminatedUnion with one variant", () => {
-    const unionParser = discriminatedUnion(
-        object({type: constant("ok")}),
-        object({type: constant("ok"), value: number}),
-    ).parser;
-
-    it("parses a valid value", () => {
-        const input = {type: "ok", value: 3};
+describe("a discriminatedUnion with no variants", () => {
+    const parseUnion = discriminatedUnionOn("shape").parser;
 
-        expect(parse(input, unionParser)).toEqual(success(input));
+    it("fails appropriately given a non-object", () => {
+        expect(parse(true, parseUnion)).toEqual(
+            failure("At (root) -- expected object, but got true"),
+        );
     });
 
-    it("rejects a value with the wrong `type`", () => {
-        const input = {type: "bad", value: 3};
-
-        expect(parse(input, unionParser)).toEqual(
-            failure(`At (root).type -- expected "ok", but got "bad"`),
+    it("fails appropriately given an object without the discriminant key", () => {
+        expect(parse({}, parseUnion)).toEqual(
+            failure(
+                "At (root).shape -- expected a valid value, but got undefined",
+            ),
         );
     });
 
-    it("rejects a value with a valid type but wrong fields", () => {
-        const input = {type: "ok", value: "foobar"};
-
-        expect(parse(input, unionParser)).toEqual(
-            failure(`At (root).value -- expected number, but got "foobar"`),
+    it("fails appropriately given an object with the discriminant key", () => {
+        expect(parse({shape: "squarle"}, parseUnion)).toEqual(
+            failure(
+                `At (root).shape -- expected a valid value, but got "squarle"`,
+            ),
         );
     });
 });
 
-describe("a discriminatedUnion with two variants", () => {
-    const unionParser = discriminatedUnion(
-        object({type: constant("rectangle")}),
-        object({type: constant("rectangle"), width: number}),
-    ).or(
-        object({type: constant("circle")}),
-        object({type: constant("circle"), radius: number}),
+describe("a discriminatedUnion with one variant", () => {
+    const parseCircle = object({shape: constant("circle"), radius: number});
+    const parseUnion = discriminatedUnionOn("shape").withBranch(
+        "circle",
+        parseCircle,
     ).parser;
 
-    it("parses a valid rectangle", () => {
-        const input = {type: "rectangle", width: 42};
-
-        expect(parse(input, unionParser)).toEqual(success(input));
+    it("fails appropriately given a non-object", () => {
+        expect(parse(true, parseUnion)).toEqual(
+            failure("At (root) -- expected object, but got true"),
+        );
     });
 
-    it("rejects a rectangle with no width", () => {
-        const input = {type: "rectangle", radius: 99};
-
-        expect(parse(input, unionParser)).toEqual(
-            failure(`At (root).width -- expected number, but got undefined`),
+    it("fails appropriately given an object without the discriminant key", () => {
+        expect(parse({}, parseUnion)).toEqual(
+            failure(
+                "At (root).shape -- expected a valid value, but got undefined",
+            ),
         );
     });
 
-    it("parses a valid circle", () => {
-        const input = {type: "circle", radius: 7};
+    it("fails appropriately given an object with an invalid discriminant", () => {
+        expect(parse({shape: "squarle"}, parseUnion)).toEqual(
+            failure(
+                `At (root).shape -- expected a valid value, but got "squarle"`,
+            ),
+        );
+    });
 
-        expect(parse(input, unionParser)).toEqual(success(input));
+    it("succeeds given a valid object", () => {
+        const input = {shape: "circle", radius: 3};
+        expect(parse(input, parseUnion)).toEqual(success(input));
     });
+});
 
-    it("rejects a circle with no radius", () => {
-        const input = {type: "circle", width: 99};
+describe("a discriminatedUnion with two variants", () => {
+    const parseCircle = object({shape: constant("circle"), radius: number});
+    const parseRectangle = object({
+        shape: constant("rectangle"),
+        width: number,
+        height: number,
+    });
+    const parseUnion = discriminatedUnionOn("shape")
+        .withBranch("circle", parseCircle)
+        .withBranch("rectangle", parseRectangle).parser;
 
-        expect(parse(input, unionParser)).toEqual(
-            failure(`At (root).radius -- expected number, but got undefined`),
+    it("fails appropriately given a non-object", () => {
+        expect(parse(true, parseUnion)).toEqual(
+            failure("At (root) -- expected object, but got true"),
         );
     });
 
-    it("rejects a value with an unrecognized `type`", () => {
-        const input = {type: "triangle", width: -1, radius: 99};
+    it("fails appropriately given an object without the discriminant key", () => {
+        expect(parse({}, parseUnion)).toEqual(
+            failure(
+                "At (root).shape -- expected a valid value, but got undefined",
+            ),
+        );
+    });
 
-        expect(parse(input, unionParser)).toEqual(
+    it("fails appropriately given an object with an invalid discriminant", () => {
+        expect(parse({shape: "squarle"}, parseUnion)).toEqual(
             failure(
-                `At (root).type -- expected "rectangle", but got "triangle"`,
+                `At (root).shape -- expected a valid value, but got "squarle"`,
             ),
         );
     });
+
+    it("successfully parses the first branch", () => {
+        const input = {shape: "circle", radius: 3};
+        expect(parse(input, parseUnion)).toEqual(success(input));
+    });
+
+    it("successfully parses the second branch", () => {
+        const input = {shape: "rectangle", width: 2, height: 4};
+        expect(parse(input, parseUnion)).toEqual(success(input));
+    });
+
+    it("doesn't try other branches after finding one that matches the discriminant key", () => {
+        const input = {shape: "circle", width: 2, height: 4};
+        expect(parse(input, parseUnion)).toEqual(
+            failure(`At (root).radius -- expected number, but got undefined`),
+        );
+    });
 });

packages/perseus/src/util/parse-perseus-json/general-purpose-parsers/discriminated-union.ts

@@ -1,46 +1,76 @@
-import {isSuccess} from "../result";
-
-import {pipeParsers} from "./pipe-parsers";
-
-import type {Parser} from "../parser-types";
-
-// discriminatedUnion() should be preferred over union() when parsing a
-// discriminated union type, because discriminatedUnion() produces more
-// understandable failure messages. It takes the discriminant as the source of
-// truth for which variant is to be parsed, and expects the other data to match
-// that variant.
-export function discriminatedUnion<T>(
-    narrow: Parser<unknown>,
-    parseVariant: Parser<T>,
-): DiscriminatedUnionBuilder<T> {
-    return new DiscriminatedUnionBuilder(
-        pipeParsers(narrow).then(parseVariant).parser,
-    );
+import {isObject} from "./is-object";
+
+import type {ParseContext, Parser} from "../parser-types";
+
+type Primitive = number | string | boolean | null | undefined;
+
+/**
+ * discriminatedUnion() should be preferred over union() when parsing a
+ * discriminated union type, because discriminatedUnion() produces more
+ * understandable failure messages. It takes the discriminant as the source of
+ * truth for which variant is to be parsed, and expects the other data to match
+ * that variant.
+ */
+export function discriminatedUnionOn<DK extends string>(discriminantKey: DK) {
+    const noMoreBranches: Parser<never> = (raw: unknown, ctx: ParseContext) => {
+        if (!isObject(raw)) {
+            return ctx.failure("object", raw);
+        }
+        return ctx
+            .forSubtree(discriminantKey)
+            .failure("a valid value", raw[discriminantKey]);
+    };
+
+    return new DiscriminatedUnionBuilder(discriminantKey, noMoreBranches);
 }
 
-class DiscriminatedUnionBuilder<Variant> {
-    constructor(public parser: Parser<Variant>) {}
+class DiscriminatedUnionBuilder<
+    DK extends string,
+    Union extends {[k in DK]: Primitive},
+> {
+    constructor(
+        private discriminantKey: DK,
+        public parser: Parser<Union>,
+    ) {}
+
+    withBranch<Variant extends {[k in DK]: Primitive}>(
+        discriminantValue: Primitive,
+        parseNewVariant: Parser<Variant>,
+    ): DiscriminatedUnionBuilder<DK, Union | Variant> {
+        const parseNewBranch = discriminatedUnionBranch(
+            this.discriminantKey,
+            discriminantValue,
+            parseNewVariant,
+            this.parser,
+        );
 
-    or<Variant2>(
-        narrow: Parser<unknown>,
-        parseVariant: Parser<Variant2>,
-    ): DiscriminatedUnionBuilder<Variant | Variant2> {
         return new DiscriminatedUnionBuilder(
-            either(narrow, parseVariant, this.parser),
+            this.discriminantKey,
+            parseNewBranch,
         );
     }
 }
 
-function either<A, B>(
-    narrowToA: Parser<unknown>,
-    parseA: Parser<A>,
-    parseB: Parser<B>,
-): Parser<A | B> {
-    return (rawValue, ctx) => {
-        if (isSuccess(narrowToA(rawValue, ctx))) {
-            return parseA(rawValue, ctx);
+function discriminatedUnionBranch<
+    DK extends string,
+    DV extends Primitive,
+    Variant extends {[k in DK]: DV},
+    Rest extends {[k in DK]: DV},
+>(
+    discriminantKey: DK,
+    discriminantValue: DV,
+    parseVariant: Parser<Variant>,
+    parseOtherBranches: Parser<Rest>,
+): Parser<Variant | Rest> {
+    return (raw: unknown, ctx: ParseContext) => {
+        if (!isObject(raw)) {
+            return ctx.failure("object", raw);
+        }
+
+        if (raw[discriminantKey] === discriminantValue) {
+            return parseVariant(raw, ctx);
         }
 
-        return parseB(rawValue, ctx);
+        return parseOtherBranches(raw, ctx);
     };
 }

packages/perseus/src/util/parse-perseus-json/general-purpose-parsers/discriminated-union.typetest.ts

@@ -0,0 +1,65 @@
+import {constant} from "./constant";
+import {discriminatedUnionOn} from "./discriminated-union";
+import {number} from "./number";
+import {object} from "./object";
+
+import type {Parser} from "../parser-types";
+
+type Figure =
+    | {shape: "circle"; radius: number}
+    | {shape: "rectangle"; width: number; height: number}
+    | {shape: "square"; sideLength: number};
+
+const parseCircle = object({
+    shape: constant("circle"),
+    radius: number,
+});
+
+const parseRectangle = object({
+    shape: constant("rectangle"),
+    width: number,
+    height: number,
+});
+
+const parseSquare = object({
+    shape: constant("square"),
+    sideLength: number,
+});
+
+// Test: parsed result is assignable to the union type
+{
+    const parser = discriminatedUnionOn("shape")
+        .withBranch("circle", parseCircle)
+        .withBranch("rectangle", parseRectangle)
+        .withBranch("square", parseSquare).parser;
+
+    parser satisfies Parser<Figure>;
+
+    // Guard against implicit 'any' type
+    // @ts-expect-error - Type '{ shape: "circle"; radius: number; }' is not assignable to type 'string'.
+    parser satisfies Parser<string>;
+}
+
+// Test: parse result with extra branches is not assignable to the union type
+{
+    const parser = discriminatedUnionOn("shape")
+        .withBranch("circle", parseCircle)
+        .withBranch("rectangle", parseRectangle)
+        .withBranch("square", parseSquare)
+        .withBranch("extra", object({shape: constant("extra")})).parser;
+
+    // @ts-expect-error - Type '{shape: "extra"}' is not assignable to type 'Figure'
+    parser satisfies Parser<Figure>;
+}
+
+// Test: each variant must contain the discriminant key
+{
+    // @ts-expect-error - property 'shape' is missing in type '{}'
+    discriminatedUnionOn("shape").withBranch("circle", object({}));
+}
+
+// Test: each variant must be an object
+{
+    // @ts-expect-error - Type 'number' is not assignable to type '{ shape: Primitive; }'.
+    discriminatedUnionOn("shape").withBranch("circle", number);
+}