Handle union types in FlattenMaps

Automattic · May 2, 2023 · 1204148 · 1204148
1 parent b864f62
commit 1204148
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/test/types/lean.test.ts b/test/types/lean.test.ts
@@ -1,4 +1,4 @@
-import { Schema, model, Document, Types, InferSchemaType, FlattenMaps } from 'mongoose';
+import { Schema, model, Types, InferSchemaType, FlattenMaps } from 'mongoose';
 import { expectAssignable, expectError, expectType } from 'tsd';
 
 function gh10345() {
@@ -179,3 +179,20 @@ async function gh13345_2() {
   expectAssignable<FlattenMaps<Place>>(place);
   expectType<Record<string, string>>(place.images[0].description);
 }
+
+async function gh13345_3() {
+  const imageSchema = new Schema({
+    url: { required: true, type: String }
+  });
+
+  const placeSchema = new Schema({
+    images: { type: [imageSchema], default: undefined }
+  });
+
+  type Place = InferSchemaType<typeof placeSchema>;
+
+  const PlaceModel = model('Place', placeSchema);
+
+  const place = await PlaceModel.findOne().lean().orFail().exec();
+  expectAssignable<Place>(place);
+}
diff --git a/types/index.d.ts b/types/index.d.ts
@@ -594,12 +594,18 @@ declare module 'mongoose' {
   export type UpdateQuery<T> = _UpdateQuery<T> & AnyObject;
 
   export type FlattenMaps<T> = {
-    [K in keyof T]: T[K] extends Map<any, infer V>
-      ? Record<string, V> : T[K] extends TreatAsPrimitives
-        ? T[K] : T[K] extends Types.DocumentArray<infer ItemType>
-          ? Types.DocumentArray<FlattenMaps<ItemType>> : FlattenMaps<T[K]>;
+    [K in keyof T]: FlattenProperty<T[K]>;
   };
 
+  /**
+   * Separate type is needed for properties of union type (for example, Types.DocumentArray | undefined) to apply conditional check to each member of it
+   * https://www.typescriptlang.org/docs/handbook/2/conditional-types.html#distributive-conditional-types
+   */
+  type FlattenProperty<T> = T extends Map<any, infer V>
+    ? Record<string, V> : T extends TreatAsPrimitives
+      ? T : T extends Types.DocumentArray<infer ItemType>
+        ? Types.DocumentArray<FlattenMaps<ItemType>> : FlattenMaps<T>;
+
   export type actualPrimitives = string | boolean | number | bigint | symbol | null | undefined;
   export type TreatAsPrimitives = actualPrimitives | NativeDate | RegExp | symbol | Error | BigInt | Types.ObjectId | Buffer | Function;