From e27afd0fc1e730d94dbd272801dbc879139cffd5 Mon Sep 17 00:00:00 2001
From: Francois Levasseur <francois.levasseur@botpress.com>
Date: Tue, 19 Nov 2024 17:24:06 -0500
Subject: [PATCH] fix(zui): handle optional unions without the need for never

---
 .../json-schema-to-zui/parsers/parseSchema.ts | 37 ++++++++++++++++++-
 .../transforms/json-schema-to-zui/types.ts    |  2 +
 zui/src/transforms/transform-pipeline.test.ts | 31 ++++++++++++++++
 3 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 zui/src/transforms/transform-pipeline.test.ts

diff --git a/zui/src/transforms/json-schema-to-zui/parsers/parseSchema.ts b/zui/src/transforms/json-schema-to-zui/parsers/parseSchema.ts
index 0e41c52a..a56793ce 100644
--- a/zui/src/transforms/json-schema-to-zui/parsers/parseSchema.ts
+++ b/zui/src/transforms/json-schema-to-zui/parsers/parseSchema.ts
@@ -17,13 +17,23 @@ import { parseNullable } from './parseNullable'
 import { parseRef } from './parseRef'
 import { ParserSelector, Refs, JsonSchemaObject, JsonSchema, Serializable, JSONSchemaExtended } from '../types'
 import { parseDiscriminator } from './parseDiscriminator'
+import { isEqual } from 'lodash'
 
 export const parseSchema = (
   schema: JSONSchemaExtended,
   refs: Refs = { seen: new Map(), path: [] },
   blockMeta?: boolean,
 ): string => {
-  if (typeof schema !== 'object') return schema ? 'z.any()' : 'z.never()'
+  if (its.anything(schema)) {
+    return 'z.any()'
+  }
+  if (its.nothing(schema)) {
+    /**
+     * Nothing in JSONSchema both means undefined and never in TypeScript
+     * We map to never for backwards compatibility
+     */
+    return 'z.never()'
+  }
 
   if (refs.parserOverride) {
     const custom = refs.parserOverride(schema, refs)
@@ -121,6 +131,29 @@ const selectParser: ParserSelector = (schema, refs) => {
   }
 }
 
+const anything = (x: JSONSchemaExtended): boolean => {
+  if (typeof x === 'boolean') {
+    return x
+  }
+  if (isEqual(x, {})) {
+    return true
+  }
+  if (x.not) {
+    return nothing(x.not)
+  }
+  return false
+}
+
+const nothing = (x: JSONSchemaExtended): boolean => {
+  if (typeof x === 'boolean') {
+    return !x
+  }
+  if (x.not) {
+    return anything(x.not)
+  }
+  return false
+}
+
 export const its = {
   an: {
     object: (x: JsonSchemaObject): x is JsonSchemaObject & { type: 'object' } => x.type === 'object',
@@ -178,4 +211,6 @@ export const its = {
     } => x.oneOf !== undefined,
     ref: (x: JsonSchemaObject): x is JsonSchemaObject & { $ref: string } => x.$ref !== undefined,
   },
+  anything,
+  nothing,
 }
diff --git a/zui/src/transforms/json-schema-to-zui/types.ts b/zui/src/transforms/json-schema-to-zui/types.ts
index 49ffec13..c6341173 100644
--- a/zui/src/transforms/json-schema-to-zui/types.ts
+++ b/zui/src/transforms/json-schema-to-zui/types.ts
@@ -51,6 +51,8 @@ export type JsonSchemaObject = {
   const?: Serializable
   enum?: Serializable[]
 
+  not?: JsonSchema
+
   errorMessage?: { [key: string]: string | undefined }
   [zuiKey]?: ZuiExtensionObject
 } & { [key: string]: any }
diff --git a/zui/src/transforms/transform-pipeline.test.ts b/zui/src/transforms/transform-pipeline.test.ts
new file mode 100644
index 00000000..d739a58c
--- /dev/null
+++ b/zui/src/transforms/transform-pipeline.test.ts
@@ -0,0 +1,31 @@
+import { zodToJsonSchema } from './zui-to-json-schema'
+import { jsonSchemaToZui } from './json-schema-to-zui'
+import { toTypescriptSchema } from './zui-to-typescript-schema'
+import { test, expect } from 'vitest'
+import z from '../z'
+import { evalZuiString } from './common/eval-zui-string'
+
+/**
+ * This test file contains integration tests to ensure the multiple transforms can be chained together
+ */
+
+const transformAll = (originalSchema: z.ZodType): z.ZodType => {
+  const jsonSchema = zodToJsonSchema(originalSchema)
+  console.log(JSON.stringify(jsonSchema))
+  const zuiSchema = jsonSchemaToZui(jsonSchema)
+  const typescriptSchema = toTypescriptSchema(zuiSchema)
+  console.log(typescriptSchema)
+  const evalResult = evalZuiString(typescriptSchema)
+  if (evalResult.sucess === false) {
+    throw new Error(`Failed to evaluate zui schema "${typescriptSchema}"; ${evalResult.error}`)
+  }
+  return evalResult.value
+}
+
+test('optional union from one end to the other and back', () => {
+  const originalSchema = z.object({
+    foo: z.literal('42').or(z.literal(42)).optional().nullable(),
+  })
+  const destinationSchema = transformAll(originalSchema)
+  expect(originalSchema.isEqual(destinationSchema)).toBe(true)
+})