fix: correctly handled mapped union types

package.json

@@ -46,6 +46,7 @@
     "debug": "tsx --inspect-brk ts-json-schema-generator.ts",
     "format": "eslint --fix",
     "lint": "eslint",
+    "prepare": "npm run build",
     "prepublishOnly": "npm run build",
     "release": "npm run build && auto shipit",
     "run": "tsx ts-json-schema-generator.ts",

src/NodeParser/MappedTypeNodeParser.ts

@@ -34,23 +34,22 @@
     }
 
     public createType(node: ts.MappedTypeNode, context: Context): BaseType {
+        // Check if the constraint is `keyof T` where T resolves to a union type.
+        // In TypeScript, mapped types distribute over unions, so `{ [P in keyof (A | B)]: ... }`
+        // is equivalent to `{ [P in keyof A]: ... } | { [P in keyof B]: ... }`.
+        const distributedType = this.tryDistributeUnion(node, context);
+        if (distributedType) {
+            return distributedType;
+        }
+
         const constraintType = this.childNodeParser.createType(node.typeParameter.constraint!, context);
         const keyListType = derefType(constraintType);
-        const id = `indexed-type-${getKey(node, context)}`;
 
-        if (keyListType instanceof UnionType) {
-            // Key type resolves to a set of known properties
-            return new ObjectType(
-                id,
-                [],
-                this.getProperties(node, keyListType, context),
-                this.getAdditionalProperties(node, keyListType, context),
-            );
-        }
+        const id = `indexed-type-${getKey(node, context)}`;
 
-        if (keyListType instanceof LiteralType) {
-            // Key type resolves to single known property
-            return new ObjectType(id, [], this.getProperties(node, new UnionType([keyListType]), context), false);
+        const objectType = this.createObjectFromKeyList(node, keyListType, id, context);
+        if (objectType) {
+            return objectType;
         }
 
         const maybeUnionType = this.childNodeParser.createType(
@@ -123,6 +122,31 @@
         return false;
     }
 
+    // Attempts to create an ObjectType from a resolved key list type.
+    // Handles UnionType (set of known property keys) and LiteralType (single known property key).
+    // Returns undefined if the key list type is not one of these.
+    protected createObjectFromKeyList(
+        node: ts.MappedTypeNode,
+        keyListType: BaseType,
+        id: string,
+        context: Context,
+    ): ObjectType | undefined {
+        if (keyListType instanceof UnionType) {
+            return new ObjectType(
+                id,
+                [],
+                this.getProperties(node, keyListType, context),
+                this.getAdditionalProperties(node, keyListType, context),
+            );
+        }
+
+        if (keyListType instanceof LiteralType) {
+            return new ObjectType(id, [], this.getProperties(node, new UnionType([keyListType]), context), false);
+        }
+
+        return undefined;
+    }
+
     protected mapKey(node: ts.MappedTypeNode, rawKey: LiteralType, context: Context): BaseType {
         if (!node.nameType) {
             return rawKey;
@@ -212,4 +236,75 @@
 
         return subContext;
     }
+
+    // Checks if the mapped type's constraint is `keyof T` where `T` is a type parameter
+    // that resolves to a union type in the current context.
+    // If so, distributes the mapped type over each union member (like TypeScript does),
+    // returning a UnionType of the individually mapped types.
+    //
+    // TypeScript distributes mapped types over union type parameters:
+    // `{ [P in keyof T]: X }` where `T = A | B` becomes `{ [P in keyof A]: X } | { [P in keyof B]: X }`
+    protected tryDistributeUnion(node: ts.MappedTypeNode, context: Context): BaseType | undefined {
+        const { constraint } = node.typeParameter;
+        if (!constraint) {
+            return undefined;
+        }
+
+        // Check if constraint is `keyof X`
+        if (!ts.isTypeOperatorNode(constraint) || constraint.operator !== ts.SyntaxKind.KeyOfKeyword) {
+            return undefined;
+        }
+
+        // Resolve the operand type (X in `keyof X`)
+        const operandType = this.childNodeParser.createType(constraint.type, context);
+        const derefedOperand = derefType(operandType);
+
+        if (!(derefedOperand instanceof UnionType)) {
+            return undefined;
+        }
+
+        const unionMembers = derefedOperand.getTypes();
+
+        // Only distribute if the union contains object-like types (not a union of literals/primitives)
+        const hasObjectTypes = unionMembers.some((member) => {
+            const derefed = derefType(member);
+            return derefed instanceof ObjectType;
+        });
+
+        if (!hasObjectTypes) {
+            return undefined;
+        }
+
+        // Distribute the mapping for each union type individually
+        const mappedTypes = unionMembers.map((member) => {
+            // Create a new context where the operand type parameter is bound to this union member
+            const subContext = new Context(node);
+
+            for (const parentParameter of context.getParameters()) {
+                const arg = context.getArgument(parentParameter);
+                // If this argument is the union we're distributing over, replace it with the member
+                if (arg === operandType || derefType(arg) === derefedOperand) {
+                    subContext.pushParameter(parentParameter);
+                    subContext.pushArgument(member);
+                } else {
+                    subContext.pushParameter(parentParameter);
+                    subContext.pushArgument(arg);
+                }
+            }
+
+            // Now resolve the constraint (keyof member) and create the mapped type for this member
+            const memberConstraintType = this.childNodeParser.createType(constraint, subContext);
+            const memberKeyListType = derefType(memberConstraintType);
+            const memberId = `indexed-type-${getKey(node, subContext)}`;
+
+            return this.createObjectFromKeyList(node, memberKeyListType, memberId, subContext)
+                ?? this.createType(node, subContext);
+        });
+
+        const result = new UnionType(mappedTypes).normalize();
+
+        // Preserve annotations (e.g., @discriminator) from the original operand type
+        // onto the resulting union, since the distribution creates a brand new UnionType.
+        return preserveAnnotation(operandType, result);
+    }
 }

src/TypeFormatter/AnnotatedTypeFormatter.ts

@@ -54,12 +54,16 @@ export class AnnotatedTypeFormatter implements SubTypeFormatter {
     }
     public getDefinition(type: AnnotatedType): Definition {
         const annotations = type.getAnnotations();
+        // Copy annotations to avoid mutating the original object, which may be shared
+        // with other AnnotatedType instances (e.g., via preserveAnnotation).
+        let restAnnotations = annotations;
 
         if ("discriminator" in annotations) {
             const deref = derefType(type.getType());
             if (deref instanceof UnionType) {
                 deref.setDiscriminator(annotations.discriminator as string);
-                delete annotations.discriminator;
+                const { discriminator: _, ...rest } = annotations;
+                restAnnotations = rest;
             } else {
                 throw new JsonTypeError(
                     `Cannot assign discriminator tag to type: ${deref.getName()}. This tag can only be assigned to union types.`,
@@ -70,7 +74,7 @@ export class AnnotatedTypeFormatter implements SubTypeFormatter {
 
         const def: Definition = {
             ...this.childTypeFormatter.getDefinition(type.getType()),
-            ...type.getAnnotations(),
+            ...restAnnotations,
         };
 
         if ("$ref" in def && "type" in def) {

src/Utils/isAssignableTo.ts

@@ -95,6 +95,9 @@ export function isAssignableTo(
     inferMap: Map<string, BaseType> = new Map(),
     insideTypes: Set<BaseType> = new Set(),
 ): boolean {
+    // Keep original source for infer map so annotations (e.g. @discriminator) are preserved
+    const originalSource = source;
+
     // Dereference source and target
     source = derefType(source);
     target = derefType(target);
@@ -115,9 +118,9 @@ export function isAssignableTo(
         const infer = inferMap.get(key);
 
         if (infer === undefined) {
-            inferMap.set(key, source);
+            inferMap.set(key, originalSource);
         } else {
-            inferMap.set(key, new UnionType([infer, source]));
+            inferMap.set(key, new UnionType([infer, originalSource]));
         }
 
         return true;

src/Utils/narrowType.ts

@@ -3,6 +3,7 @@ import { EnumType } from "../Type/EnumType.js";
 import { NeverType } from "../Type/NeverType.js";
 import { UnionType } from "../Type/UnionType.js";
 import { derefType } from "./derefType.js";
+import { preserveAnnotation } from "./preserveAnnotation.js";
 
 /**
  * Narrows the given type by passing all variants to the given predicate function. So when type is a union type then
@@ -46,7 +47,7 @@ export function narrowType(type: BaseType, predicate: (type: BaseType) => boolea
             } else if (types.length === 1) {
                 return types[0];
             } else {
-                return new UnionType(types);
+                return preserveAnnotation(type, new UnionType(types));
             }
         }
         return type;

src/Utils/preserveAnnotation.ts

@@ -1,15 +1,24 @@
 import type { BaseType } from "../Type/BaseType.js";
 import { AnnotatedType } from "../Type/AnnotatedType.js";
+import { AliasType } from "../Type/AliasType.js";
+import { DefinitionType } from "../Type/DefinitionType.js";
+import { ReferenceType } from "../Type/ReferenceType.js";
 
 /**
  * Return the new type wrapped in an annotated type with the same annotations as the original type.
- * @param originalType The original type. If this is an annotated type,
- *      then the returned type will be wrapped with the same annotations.
+ * @param originalType The original type. If this is an annotated type (possibly wrapped in AliasType,
+ *      DefinitionType, or ReferenceType), then the returned type will be wrapped with the same annotations.
  * @param newType The type to be wrapped.
  */
 export function preserveAnnotation(originalType: BaseType, newType: BaseType): BaseType {
     if (originalType instanceof AnnotatedType) {
         return new AnnotatedType(newType, originalType.getAnnotations(), originalType.isNullable());
     }
+    if (originalType instanceof AliasType || originalType instanceof DefinitionType) {
+        return preserveAnnotation(originalType.getType(), newType);
+    }
+    if (originalType instanceof ReferenceType && originalType.hasType()) {
+        return preserveAnnotation(originalType.getType(), newType);
+    }
     return newType;
 }

test/valid-data/type-indexed-access-type-union/schema.json

@@ -3,16 +3,26 @@
   "$schema": "http://json-schema.org/draft-07/schema#",
   "definitions": {
     "MyType": {
-      "additionalProperties": false,
-      "properties": {
-        "foo": {
-          "type": [
-            "number",
-            "boolean"
-          ]
+      "anyOf": [
+        {
+          "additionalProperties": false,
+          "properties": {
+            "foo": {
+              "type": "number"
+            }
+          },
+          "type": "object"
+        },
+        {
+          "additionalProperties": false,
+          "properties": {
+            "foo": {
+              "type": "boolean"
+            }
+          },
+          "type": "object"
         }
-      },
-      "type": "object"
+      ]
     }
   }
 }

test/valid-data/type-mapped-union-discriminator-shared-annotation/index.test.ts

@@ -0,0 +1,7 @@
+import { assertValidSchema } from "../../utils";
+import { test } from "node:test";
+
+test(
+    "valid-data - type-mapped-union-discriminator-shared-annotation",
+    assertValidSchema("type-mapped-union-discriminator-shared-annotation", "*", { jsDoc: "basic", discriminatorType: "open-api" }),
+);

test/valid-data/type-mapped-union-discriminator-shared-annotation/main.ts

@@ -0,0 +1,34 @@
+/**
+ * Regression test: when a @discriminator-annotated union goes through a
+ * mapped type and is referenced from TWO different exported types, the
+ * annotations object produced by preserveAnnotation is shared between
+ * the resulting AnnotatedType instances.
+ *
+ * Before the fix, `delete annotations.discriminator` in
+ * AnnotatedTypeFormatter mutated the shared object, so whichever type
+ * was processed first would steal the discriminator from the second.
+ */
+
+type Fish = {
+    animal_type: "fish";
+    found_in: "ocean" | "river";
+};
+
+type Bird = {
+    animal_type: "bird";
+    can_fly: boolean;
+};
+
+/** @discriminator animal_type */
+type Animal = Fish | Bird;
+
+// A simple recursive mapped type that preserves structure but forces
+// the union through tryDistributeUnion → preserveAnnotation.
+type DeepMapped<T extends object> = {
+    [P in keyof T]: T[P] extends object ? DeepMapped<T[P]> : T[P];
+};
+
+// Two distinct exported types that both embed the same discriminated union
+// through the same mapped type. They will share the annotations object.
+export type First = DeepMapped<{ pet: Animal; label: string }>;
+export type Second = DeepMapped<{ pet: Animal; tag: number }>;