From 7b61de45625451d1772d3538e02c5c5fcf158c8b Mon Sep 17 00:00:00 2001
From: exoego <ytatsuno.jp@gmail.com>
Date: Tue, 16 Apr 2024 09:30:42 +0900
Subject: [PATCH] Add or operation

---
 src/other.ts        | 45 +++++++++++++++++++++++++++++++++++++++++++++
 tests/index.test.ts | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/src/other.ts b/src/other.ts
index 3e3c0ea..69c6324 100644
--- a/src/other.ts
+++ b/src/other.ts
@@ -78,3 +78,48 @@ export const optional = (innerValidator?: Validator<any, any>) => {
     },
   };
 };
+
+interface OrSchema extends Fn {
+  return: this["args"] extends [...any, infer prev, infer last]
+    ? last extends { $outputType: infer T1 }
+      ? prev extends { $outputType: infer T2 }
+        ? T1 | T2
+        : T1 | prev
+      : prev extends { $outputType: infer T2 }
+        ? last | T2
+        : last | prev
+    : never;
+}
+
+export const or = <T1, T2>(innerValidator: Validator<T1, any>) => {
+  const ctx = {
+    chain: null as Validator<T2, any> | null,
+  } satisfies { chain: Validator<T2, any> | null };
+
+  return {
+    name: "or" as const,
+    $inputType: "any" as unknown as T1 | T2,
+    $outputType: "any" as unknown as OrSchema,
+    processChain: (chain: Validator<T2, any> | null) => {
+      if (chain !== null) {
+        ctx.chain = chain;
+      }
+      return [];
+    },
+    parse: (arg: any) => {
+      if (arg === null) return arg;
+      const chain = ctx.chain;
+      if (chain === null) {
+        throw new Error(
+          `No inner validator. Make sure to do c.someValidator().or(c.otherValidator())`,
+        );
+      }
+
+      try {
+        return chain.parse(arg);
+      } catch (e) {
+        return innerValidator.parse(arg);
+      }
+    },
+  };
+};
diff --git a/tests/index.test.ts b/tests/index.test.ts
index e017c9b..fa1746e 100644
--- a/tests/index.test.ts
+++ b/tests/index.test.ts
@@ -9,6 +9,7 @@ import {
   object,
   nullable,
   optional,
+  or,
   coerce,
 } from "../src";
 import { Equal, Expect } from "./helpers.types";
@@ -43,6 +44,7 @@ describe("basic tests", () => {
     nullable,
     literal,
     optional,
+    or,
     set,
     setNonEmpty,
     setMin,
@@ -238,4 +240,43 @@ describe("basic tests", () => {
 
     // TODO: test other coerce functions
   });
+
+  describe("or", () => {
+    test("basic", () => {
+      const schema = c.string().or(c.number());
+      expect(schema.parse("yay")).toEqual("yay");
+      expect(schema.parse(42)).toEqual(42);
+      expect(() => schema.parse({ x: 42 })).toThrow();
+
+      type SchemaType = Infer<typeof schema>;
+      type _test = Expect<Equal<SchemaType, string | number>>;
+    });
+
+    test("literal", () => {
+      const schema = c.literal("a").or(c.literal(42));
+      expect(schema.parse("a")).toEqual("a");
+      expect(schema.parse(42)).toEqual(42);
+      expect(() => schema.parse("aaa")).toThrow();
+
+      type SchemaType = Infer<typeof schema>;
+      type _test = Expect<Equal<SchemaType, "a" | 42>>;
+    });
+
+    test("object", () => {
+      const schema = c.object({ x: number() }).or(c.object({ y: string() }));
+      // FIXME: .or should be chainable multiple times
+      // .or(c.string());
+      expect(schema.parse({ x: 42 })).toEqual({ x: 42 });
+      expect(schema.parse({ y: "foo" })).toEqual({ y: "foo" });
+      // expect(schema.parse("foo")).toEqual("foo");
+      expect(() => schema.parse({ x: "bar" })).toThrow();
+      expect(() => schema.parse({ y: 42 })).toThrow();
+
+      // strip keys in other schema
+      expect(schema.parse({ x: 42, y: "foo" })).toEqual({ x: 42 });
+
+      type SchemaType = Infer<typeof schema>;
+      type _test = Expect<Equal<SchemaType, { x: number } | { y: string }>>;
+    });
+  });
 });