Merge pull request #14632 from Automattic/vkarpov15/gh-14615

types(models+query): infer return type from schema for 1-level deep nested paths
Automattic · Jun 4, 2024 · 6183560 · 6183560
2 parents c71ba5e + 2c45e85
commit 6183560
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 3 deletions.
diff --git a/test/types/models.test.ts b/test/types/models.test.ts
@@ -13,7 +13,9 @@ import mongoose, {
   Query,
   UpdateWriteOpResult,
   AggregateOptions,
-  StringSchemaDefinition
+  WithLevel1NestedPaths,
+  NestedPaths,
+  InferSchemaType
 } from 'mongoose';
 import { expectAssignable, expectError, expectType } from 'tsd';
 import { AutoTypedSchemaType, autoTypedSchema } from './schema.test';
@@ -914,3 +916,64 @@ async function gh14440() {
     }
   ]);
 }
+
+async function gh12064() {
+  const FooSchema = new Schema({
+    one: { type: String }
+  });
+
+  const MyRecordSchema = new Schema({
+    _id: { type: String },
+    foo: { type: FooSchema },
+    arr: [Number]
+  });
+
+  const MyRecord = model('MyRecord', MyRecordSchema);
+
+  expectType<(string | null)[]>(
+    await MyRecord.distinct('foo.one').exec()
+  );
+  expectType<(string | null)[]>(
+    await MyRecord.find().distinct('foo.one').exec()
+  );
+  expectType<unknown[]>(await MyRecord.distinct('foo.two').exec());
+  expectType<unknown[]>(await MyRecord.distinct('arr.0').exec());
+}
+
+function testWithLevel1NestedPaths() {
+  type Test1 = WithLevel1NestedPaths<{
+    topLevel: number,
+    nested1Level: {
+      l2: string
+    },
+    nested2Level: {
+      l2: { l3: boolean }
+    }
+  }>;
+
+  expectType<{
+    topLevel: number,
+    nested1Level: { l2: string },
+    'nested1Level.l2': string,
+    nested2Level: { l2: { l3: boolean } },
+    'nested2Level.l2': { l3: boolean }
+  }>({} as Test1);
+
+  const FooSchema = new Schema({
+    one: { type: String }
+  });
+
+  const schema = new Schema({
+    _id: { type: String },
+    foo: { type: FooSchema }
+  });
+
+  type InferredDocType = InferSchemaType<typeof schema>;
+
+  type Test2 = WithLevel1NestedPaths<InferredDocType>;
+  expectAssignable<{
+    _id: string | null | undefined,
+    foo?: { one?: string | null | undefined } | null | undefined,
+    'foo.one': string | null | undefined
+  }>({} as Test2);
+}
diff --git a/types/models.d.ts b/types/models.d.ts
@@ -623,7 +623,11 @@ declare module 'mongoose' {
       field: DocKey,
       filter?: FilterQuery<TRawDocType>
     ): QueryWithHelpers<
-      Array<DocKey extends keyof TRawDocType ? Unpacked<TRawDocType[DocKey]> : ResultType>,
+      Array<
+        DocKey extends keyof WithLevel1NestedPaths<TRawDocType>
+          ? WithoutUndefined<Unpacked<WithLevel1NestedPaths<TRawDocType>[DocKey]>>
+          : ResultType
+      >,
       THydratedDocumentType,
       TQueryHelpers,
       TRawDocType,

diff --git a/types/query.d.ts b/types/query.d.ts
@@ -355,7 +355,18 @@ declare module 'mongoose' {
     distinct<DocKey extends string, ResultType = unknown>(
       field: DocKey,
       filter?: FilterQuery<RawDocType>
-    ): QueryWithHelpers<Array<DocKey extends keyof DocType ? Unpacked<DocType[DocKey]> : ResultType>, DocType, THelpers, RawDocType, 'distinct', TInstanceMethods>;
+    ): QueryWithHelpers<
+      Array<
+        DocKey extends keyof WithLevel1NestedPaths<DocType>
+          ? WithoutUndefined<Unpacked<WithLevel1NestedPaths<DocType>[DocKey]>>
+          : ResultType
+      >,
+      DocType,
+      THelpers,
+      RawDocType,
+      'distinct',
+      TInstanceMethods
+    >;
 
     /** Specifies a `$elemMatch` query condition. When called with one argument, the most recent path passed to `where()` is used. */
     elemMatch<K = string>(path: K, val: any): this;

diff --git a/types/utility.d.ts b/types/utility.d.ts
@@ -2,6 +2,26 @@ declare module 'mongoose' {
   type IfAny<IFTYPE, THENTYPE, ELSETYPE = IFTYPE> = 0 extends (1 & IFTYPE) ? THENTYPE : ELSETYPE;
   type IfUnknown<IFTYPE, THENTYPE> = unknown extends IFTYPE ? THENTYPE : IFTYPE;
 
+  type WithLevel1NestedPaths<T, K extends keyof T = keyof T> = {
+    [P in K | NestedPaths<Required<T>, K>]: P extends K
+      ? T[P]
+      : P extends `${infer Key}.${infer Rest}`
+        ? Key extends keyof T
+          ? Rest extends keyof NonNullable<T[Key]>
+            ? NonNullable<T[Key]>[Rest]
+            : never
+          : never
+        : never;
+  };
+
+  type NestedPaths<T, K extends keyof T> = K extends string
+    ? T[K] extends Record<string, any> | null | undefined
+      ? `${K}.${keyof NonNullable<T[K]> & string}`
+      : never
+    : never;
+
+  type WithoutUndefined<T> = T extends undefined ? never : T;
+
   /**
     * @summary Removes keys from a type
     * @description It helps to exclude keys from a type