From 45119280b3ea181f73bf422edf50c6ebd00d8e9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Artur=20M=C3=BCller?= <me@arturmuller.com>
Date: Sat, 16 Dec 2023 20:25:45 +0100
Subject: [PATCH] Retain struct's type in partial and pick helpers (#1149)

* Retain struct's type in partial and pick helpers

* Refactor for clarity

Refactors partial helper to more directly reflect its internal logic
including schema shorthand usage.
---
 docs/reference/utilities.md           |  4 ++--
 src/structs/utilities.ts              | 20 +++++++++++++++-----
 test/validation/partial/valid-type.ts | 18 ++++++++++++++++++
 test/validation/pick/valid-type.ts    | 19 +++++++++++++++++++
 4 files changed, 54 insertions(+), 7 deletions(-)
 create mode 100644 test/validation/partial/valid-type.ts
 create mode 100644 test/validation/pick/valid-type.ts

diff --git a/docs/reference/utilities.md b/docs/reference/utilities.md
index 5dc2a3e0..c572eb8c 100644
--- a/docs/reference/utilities.md
+++ b/docs/reference/utilities.md
@@ -92,7 +92,7 @@ partial(
 { name: 'Jane' }
 ```
 
-`partial` allows you to create a new struct based on an existing object struct, but with all of its properties being optional.
+`partial` allows you to create a new struct based on an existing `object` or `type` struct, but with all of its properties being optional.
 
 ### `pick`
 
@@ -106,4 +106,4 @@ pick(
 )
 ```
 
-`pick` allows you to create a new struct based on an existing object struct, but only including specific properties.
+`pick` allows you to create a new struct based on an existing `object` or `type` struct, but only including specific properties.
diff --git a/src/structs/utilities.ts b/src/structs/utilities.ts
index 28c1c53b..d2162a3d 100644
--- a/src/structs/utilities.ts
+++ b/src/structs/utilities.ts
@@ -1,6 +1,6 @@
-import { Struct, Context, Validator } from '../struct'
+import { Context, Struct, Validator } from '../struct'
+import { Assign, ObjectSchema, ObjectType, PartialObjectSchema } from '../utils'
 import { object, optional, type } from './types'
-import { ObjectSchema, Assign, ObjectType, PartialObjectSchema } from '../utils'
 
 /**
  * Create a new struct that combines the properties properties from multiple
@@ -192,13 +192,17 @@ export function omit<S extends ObjectSchema, K extends keyof S>(
 export function partial<S extends ObjectSchema>(
   struct: Struct<ObjectType<S>, S> | S
 ): Struct<ObjectType<PartialObjectSchema<S>>, PartialObjectSchema<S>> {
-  const schema: any =
-    struct instanceof Struct ? { ...struct.schema } : { ...struct }
+  const isStruct = struct instanceof Struct
+  const schema: any = isStruct ? { ...struct.schema } : { ...struct }
 
   for (const key in schema) {
     schema[key] = optional(schema[key])
   }
 
+  if (isStruct && struct.type === 'type') {
+    return type(schema) as any
+  }
+
   return object(schema) as any
 }
 
@@ -220,7 +224,13 @@ export function pick<S extends ObjectSchema, K extends keyof S>(
     subschema[key] = schema[key]
   }
 
-  return object(subschema as Pick<S, K>)
+  switch (struct.type) {
+    case 'type':
+      return type(subschema) as any
+
+    default:
+      return object(subschema) as any
+  }
 }
 
 /**
diff --git a/test/validation/partial/valid-type.ts b/test/validation/partial/valid-type.ts
new file mode 100644
index 00000000..51411a70
--- /dev/null
+++ b/test/validation/partial/valid-type.ts
@@ -0,0 +1,18 @@
+import { number, partial, string, type } from '../../../src'
+
+export const Struct = partial(
+  type({
+    name: string(),
+    age: number(),
+  })
+)
+
+export const data = {
+  name: 'john',
+  unknownProperty: true,
+}
+
+export const output = {
+  name: 'john',
+  unknownProperty: true,
+}
diff --git a/test/validation/pick/valid-type.ts b/test/validation/pick/valid-type.ts
new file mode 100644
index 00000000..28f1a330
--- /dev/null
+++ b/test/validation/pick/valid-type.ts
@@ -0,0 +1,19 @@
+import { number, pick, string, type } from '../../../src'
+
+export const Struct = pick(
+  type({
+    name: string(),
+    age: number(),
+  }),
+  ['name']
+)
+
+export const data = {
+  name: 'john',
+  unknownProperty: true,
+}
+
+export const output = {
+  name: 'john',
+  unknownProperty: true,
+}