From 821953983a2850184efb7c917f82b02ee6923f7a Mon Sep 17 00:00:00 2001
From: Michael Hayes <michael@hayes.io>
Date: Sat, 10 Aug 2024 18:55:10 -0400
Subject: [PATCH] Add isOneOf typing on builder.inputObject

---
 .changeset/smooth-parents-punch.md             |  5 +++++
 packages/core/src/builder.ts                   | 16 +++++++++++++---
 packages/core/src/types/builder-options.ts     | 14 ++++++++++++++
 packages/core/src/types/utils.ts               |  2 ++
 .../tests/__snapshots__/index.test.ts.snap     |  6 ++++++
 packages/core/tests/examples/random-stuff.ts   | 18 ++++++++++++++++++
 6 files changed, 58 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/smooth-parents-punch.md

diff --git a/.changeset/smooth-parents-punch.md b/.changeset/smooth-parents-punch.md
new file mode 100644
index 000000000..7b24f6383
--- /dev/null
+++ b/.changeset/smooth-parents-punch.md
@@ -0,0 +1,5 @@
+---
+"@pothos/core": minor
+---
+
+Add proper typeing for isOneOf option on builder.inputType
diff --git a/packages/core/src/builder.ts b/packages/core/src/builder.ts
index 43330f1d1..8a86ab686 100644
--- a/packages/core/src/builder.ts
+++ b/packages/core/src/builder.ts
@@ -55,6 +55,7 @@ import type {
   ObjectFieldsShape,
   ObjectParam,
   ObjectTypeOptions,
+  OneOfInputShapeFromFields,
   OutputShape,
   ParentShape,
   PluginConstructorMap,
@@ -583,10 +584,16 @@ export class SchemaBuilder<Types extends SchemaTypes> {
       : Param extends keyof Types['Inputs']
         ? InputFieldsFromShape<Types, InputShape<Types, Param> & object, 'InputObject'>
         : InputFieldMap,
+    IsOneOf extends boolean = boolean,
   >(
     param: Param,
-    options: PothosSchemaTypes.InputObjectTypeOptions<Types, Fields>,
-  ): PothosSchemaTypes.InputObjectRef<Types, InputShapeFromFields<Fields>> {
+    options: PothosSchemaTypes.InputObjectTypeOptions<Types, Fields> & {
+      isOneOf?: IsOneOf;
+    },
+  ): PothosSchemaTypes.InputObjectRef<
+    Types,
+    [IsOneOf] extends [true] ? OneOfInputShapeFromFields<Fields> : InputShapeFromFields<Fields>
+  > {
     verifyRef(param);
     const name = typeof param === 'string' ? param : (param as { name: string }).name;
 
@@ -594,7 +601,10 @@ export class SchemaBuilder<Types extends SchemaTypes> {
       typeof param === 'string'
         ? new InputObjectRef<Types, InputShapeFromFields<Fields>>(name)
         : param
-    ) as PothosSchemaTypes.InputObjectRef<Types, InputShapeFromFields<Fields>>;
+    ) as PothosSchemaTypes.InputObjectRef<
+      Types,
+      [IsOneOf] extends [true] ? OneOfInputShapeFromFields<Fields> : InputShapeFromFields<Fields>
+    >;
 
     ref.updateConfig({
       kind: 'InputObject',
diff --git a/packages/core/src/types/builder-options.ts b/packages/core/src/types/builder-options.ts
index 443c57caf..d1f1cf618 100644
--- a/packages/core/src/types/builder-options.ts
+++ b/packages/core/src/types/builder-options.ts
@@ -24,6 +24,7 @@ import type {
   Normalize,
   NormalizeNullableFields,
   RemoveNeverKeys,
+  Simplify,
 } from './utils';
 
 export type AddVersionedDefaultsToBuilderOptions<
@@ -211,6 +212,19 @@ export type ValidateInterfaces<
     : 'Object shape must extend interface shape'
   : never;
 
+export type OneOfInputShapeFromFields<Fields extends InputFieldMap> =
+  keyof Fields extends infer K extends keyof Fields
+    ? K extends unknown
+      ? Simplify<
+          {
+            [Name in K]: NonNullable<InputShapeFromField<Fields[K]>>;
+          } & {
+            [Name in keyof Fields as Name extends K ? never : Name]?: never;
+          }
+        >
+      : never
+    : never;
+
 export type InputShapeFromFields<Fields extends InputFieldMap> = NormalizeNullableFields<{
   [K in string & keyof Fields]: InputShapeFromField<Fields[K]>;
 }>;
diff --git a/packages/core/src/types/utils.ts b/packages/core/src/types/utils.ts
index 8afb50f4a..473a53d29 100644
--- a/packages/core/src/types/utils.ts
+++ b/packages/core/src/types/utils.ts
@@ -138,3 +138,5 @@ export interface PartialResolveInfo {
   variableValues: GraphQLResolveInfo['variableValues'];
   schema: GraphQLResolveInfo['schema'];
 }
+
+export type Simplify<T> = { [KeyType in keyof T]: T[KeyType] } & {};
diff --git a/packages/core/tests/__snapshots__/index.test.ts.snap b/packages/core/tests/__snapshots__/index.test.ts.snap
index db24a973e..c97b0eebd 100644
--- a/packages/core/tests/__snapshots__/index.test.ts.snap
+++ b/packages/core/tests/__snapshots__/index.test.ts.snap
@@ -66,9 +66,15 @@ input NestedListInput {
   list: [[Int!]]!
 }
 
+input OneOfExample @oneOf {
+  a: String
+  b: String
+}
+
 type Query {
   constructor: Boolean
   nestedLists(input: [[[String!]!]]!, nestedListInput: NestedListInput): [[[String!]!]]
+  oneOf(oneOf: OneOfExample!): String!
   sheep: Sheep
   stuff: [stuff!]
   user: User
diff --git a/packages/core/tests/examples/random-stuff.ts b/packages/core/tests/examples/random-stuff.ts
index 5c53a684a..f5db9b2e5 100644
--- a/packages/core/tests/examples/random-stuff.ts
+++ b/packages/core/tests/examples/random-stuff.ts
@@ -346,6 +346,24 @@ builder.queryType({
         shaved: true,
       }),
     }),
+    oneOf: t.string({
+      args: {
+        oneOf: t.arg({
+          required: true,
+          type: builder.inputType('OneOfExample', {
+            isOneOf: true,
+            fields: (t) => ({
+              a: t.string(),
+              b: t.string(),
+            }),
+          }),
+        }),
+      },
+      nullable: false,
+      resolve: (_, { oneOf }) => {
+        return oneOf.a ?? oneOf.b;
+      },
+    }),
   }),
 });