From 6956ee788369dff00e5ecadb506726af3598a87e Mon Sep 17 00:00:00 2001
From: jquense <jquense@ramp.com>
Date: Fri, 29 Sep 2023 12:16:29 -0400
Subject: [PATCH] fix: pick and omit with excluded edges

fixes #2097
---
 src/object.ts  | 16 +++++++++++-----
 test/object.ts | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/object.ts b/src/object.ts
index aaabfc61f..39e921cc0 100644
--- a/src/object.ts
+++ b/src/object.ts
@@ -419,17 +419,23 @@ export default class ObjectSchema<
       if (this.fields[key]) picked[key] = this.fields[key];
     }
 
-    return this.setFields<{ [K in TKey]: TIn[K] }, TDefault>(picked);
+    return this.setFields<{ [K in TKey]: TIn[K] }, TDefault>(
+      picked,
+      this._excludedEdges.filter(
+        ([a, b]) => keys.includes(a as TKey) && keys.includes(b as TKey),
+      ),
+    );
   }
 
   omit<TKey extends keyof TIn>(keys: readonly TKey[]) {
-    const fields = { ...this.fields };
+    const remaining: TKey[] = [];
 
-    for (const key of keys) {
-      delete fields[key];
+    for (const key of Object.keys(this.fields) as TKey[]) {
+      if (keys.includes(key)) continue;
+      remaining.push(key);
     }
 
-    return this.setFields<Omit<TIn, TKey>, TDefault>(fields);
+    return this.pick(remaining);
   }
 
   from(from: string, to: keyof TIn, alias?: boolean) {
diff --git a/test/object.ts b/test/object.ts
index 2911bf60e..553c679e9 100644
--- a/test/object.ts
+++ b/test/object.ts
@@ -1121,4 +1121,48 @@ describe('Object types', () => {
       await inst.omit(['age', 'name']).validate({ color: 'mauve' }),
     ).toEqual({ color: 'mauve' });
   });
+
+  it('should pick and omit with excluded edges', async () => {
+    const inst = object().shape(
+      {
+        a1: string().when('a2', {
+          is: undefined,
+          then: (schema) => schema.required(),
+        }),
+        a2: string().when('a1', {
+          is: undefined,
+          then: (schema) => schema.required(),
+        }),
+        a3: string().required(),
+      },
+      [['a1', 'a2']],
+    );
+
+    expect(
+      inst.pick(['a1', 'a2']).isValid({
+        a1: undefined,
+        a2: 'over9000',
+      }),
+    ).resolves.toEqual(true);
+
+    expect(
+      inst.pick(['a1', 'a3']).isValid({
+        a1: 'required',
+        a3: 'asfasf',
+      }),
+    ).resolves.toEqual(true);
+
+    expect(
+      inst.omit(['a1', 'a2']).isValid({
+        a3: 'asfasf',
+      }),
+    ).resolves.toEqual(true);
+
+    expect(
+      inst.omit(['a1']).isValid({
+        a1: undefined,
+        a3: 'asfasf',
+      }),
+    ).resolves.toEqual(false);
+  });
 });