@atlaskit/forge-react-types 0.42.7 → 0.42.9

package/scripts/codegen/codeGenerator.ts

@@ -5,10 +5,21 @@ import {
 	type SourceFile,
 	type TypeAliasDeclaration,
 	type ImportDeclaration,
+	type TypeReferenceNode,
+	SyntaxKind,
 } from 'ts-morph';
 // eslint-disable-next-line import/no-extraneous-dependencies
 import kebabCase from 'lodash/kebabCase';
-import { serializeSymbolType } from './typeSerializer';
+import {
+	serializeTypeReferenceWithPickType,
+	extractPickKeys,
+	extractOmitKeys,
+} from './typeSerializer';
+import {
+	findTypeReferenceFromUnionOrIntersect,
+	getTypeNodeFromSymbol,
+	makePickOrOmitPredicate,
+} from './utils';
 
 const getNames = (symbol: Symbol) => {
 	const name = symbol.getName();
@@ -200,8 +211,11 @@ class SimpleImportDeclaration implements IImportDeclaration {
 
 	private namedImports = new Set<string>();
 
-	constructor(packageName: string) {
+	private defaultImport?: string;
+
+	constructor(packageName: string, defaultImport?: string) {
 		this.packageName = packageName;
+		this.defaultImport = defaultImport;
 	}
 
 	public addNamedImport(namedImport: string) {
@@ -213,14 +227,16 @@ class SimpleImportDeclaration implements IImportDeclaration {
 	}
 
 	public getText() {
-		if (this.namedImports.size === 0) {
+		if (this.namedImports.size === 0 && !this.defaultImport) {
 			return `import '${this.packageName}';`;
 		}
-		const importedNames = Array.from(this.namedImports)
+		const importedNamesList = Array.from(this.namedImports)
 			.sort()
-			.map((name) => `type ${name}`)
-			.join(', ');
-		return `import { ${importedNames} } from '${this.packageName}';`;
+			.map((name) => `type ${name}`);
+		const importedNames =
+			importedNamesList.length > 0 ? `{ ${importedNamesList.join(', ')} }` : null;
+		const defaultImport = this.defaultImport ? `type ${this.defaultImport}` : null;
+		return `import ${[defaultImport, importedNames].filter(Boolean).join(', ')} from '${this.packageName}';`;
 	}
 }
 
@@ -559,15 +575,24 @@ const registeredExternalTypes: Record<
 	string,
 	{
 		package: string;
-		alias?: string;
+		defaultImport: string;
 	}
 > = {
-	'React.ReactNode': {
+	'^React\..+$': {
 		package: 'react',
-		alias: 'ReactNode',
+		defaultImport: 'React',
 	},
 };
 
+const findExternalTypeInfo = (typeName: string) => {
+	return (
+		Object.entries(registeredExternalTypes).find(
+			([externalTypePattern]) =>
+				externalTypePattern === typeName || new RegExp(externalTypePattern).test(typeName),
+		)?.[1] ?? null
+	);
+};
+
 // consolidate external types into import declarations
 const consolidateImportDeclarations = (
 	importDeclarations: ImportDeclarationProxy[],
@@ -575,7 +600,7 @@ const consolidateImportDeclarations = (
 ): IImportDeclaration[] => {
 	const declarations: IImportDeclaration[] = [...importDeclarations];
 	externalTypes.forEach((typeName) => {
-		const typePackage = registeredExternalTypes[typeName];
+		const typePackage = findExternalTypeInfo(typeName);
 		if (typePackage) {
 			const existingImport = importDeclarations.find(
 				(declaration) => declaration.getBasePackage() === typePackage.package,
@@ -583,8 +608,13 @@ const consolidateImportDeclarations = (
 			if (existingImport) {
 				existingImport.addNamedImport(typeName);
 			} else {
-				const newImport = new SimpleImportDeclaration(typePackage.package);
-				newImport.addNamedImport(typePackage.alias ?? typeName);
+				const newImport = new SimpleImportDeclaration(
+					typePackage.package,
+					typePackage.defaultImport,
+				);
+				if (!typeName.startsWith(`${typePackage.defaultImport}.`)) {
+					newImport.addNamedImport(typeName);
+				}
 				declarations.push(newImport);
 			}
 		}
@@ -592,6 +622,62 @@ const consolidateImportDeclarations = (
 	return declarations;
 };
 
+const extractPlatformPropsTypeDeclarationAndTargetPropertyKeys = (
+	rawDependentTypeDeclarations: TypeAliasDeclaration[],
+	baseComponentPropsSymbol: Symbol,
+): {
+	typeDeclaration: TypeAliasDeclaration;
+	baseComponentPropTypeReference: TypeReferenceNode;
+	omitKeys: string[];
+	pickKeys: string[];
+} => {
+	// this pattern is used when there is special JSDoc comments override required to customize component documentation.
+	// e.g. Badge component has:
+	//    PlatformBadgeProps = Omit<_PlatformBadgeProps, 'children'> & {}
+	//    BadgeProps = Pick<PlatformBadgeProps, 'appearance' | 'children' | 'max' | 'testId'>
+	const specialPlatformPropsTypeDeclaration = rawDependentTypeDeclarations.find((declaration) =>
+		declaration.getName().startsWith('_Platform'),
+	);
+	const mainPlatformPropsTypeDeclaration = rawDependentTypeDeclarations.find((declaration) =>
+		declaration.getName().startsWith('Platform'),
+	);
+	if (!mainPlatformPropsTypeDeclaration && !specialPlatformPropsTypeDeclaration) {
+		throw new Error(
+			'Could not find Platform props type declaration from the component prop type source code',
+		);
+	}
+	const platformPropsTypeDeclaration =
+		specialPlatformPropsTypeDeclaration ?? mainPlatformPropsTypeDeclaration;
+
+	const typeWhichPicksPlatformProps = findTypeReferenceFromUnionOrIntersect(
+		getTypeNodeFromSymbol(baseComponentPropsSymbol!)!,
+		makePickOrOmitPredicate('Pick', 'Platform'),
+	);
+	if (!typeWhichPicksPlatformProps) {
+		throw new Error(
+			`Could not find type which picks platform props from the component prop type source code for ${baseComponentPropsSymbol.getName()}`,
+		);
+	}
+	const pickKeys = extractPickKeys(typeWhichPicksPlatformProps);
+	let omitKeys: string[] = [];
+	if (specialPlatformPropsTypeDeclaration && mainPlatformPropsTypeDeclaration) {
+		const omitNode = mainPlatformPropsTypeDeclaration
+			.getTypeNodeOrThrow()
+			.asKindOrThrow(SyntaxKind.IntersectionType)
+			.getTypeNodes()[0]
+			.asKindOrThrow(SyntaxKind.TypeReference);
+		// if there is a special platform props type declaration, we need
+		omitKeys = extractOmitKeys(omitNode);
+	}
+
+	return {
+		typeDeclaration: platformPropsTypeDeclaration!,
+		baseComponentPropTypeReference: typeWhichPicksPlatformProps,
+		pickKeys,
+		omitKeys,
+	};
+};
+
 /**
  * This function implements the new code generation logic for the component prop types.
  * Instead of referencing to the ADS component prop types, it generates the prop types
@@ -609,20 +695,20 @@ const generateComponentPropTypeSourceCodeWithSerializedType = (
 	// 2) from the prop type code further extract other relevant types in the source file,
 	//    and separate the platform props type declaration from the rest of the dependent types.
 	//    as this will be used to generate the type code.
-	const [dependentTypeDeclarations, platformPropsTypeDeclaration] = getDependentTypeDeclarations(
+	const rawDependentTypeDeclarations = getDependentTypeDeclarations(
 		baseComponentPropSymbol,
 		sourceFile,
-	).reduce(
-		(agg, declarations) => {
-			if (declarations.getName().startsWith('Platform')) {
-				// this is the platform props type declaration, we will use it to generate the type code
-				agg[1] = declarations;
-			} else {
-				agg[0].push(declarations);
-			}
-			return agg;
-		},
-		[[] as TypeAliasDeclaration[], null as TypeAliasDeclaration | null],
+	);
+	const {
+		omitKeys,
+		typeDeclaration: platformPropsTypeDeclaration,
+		baseComponentPropTypeReference,
+	} = extractPlatformPropsTypeDeclarationAndTargetPropertyKeys(
+		rawDependentTypeDeclarations,
+		baseComponentPropSymbol,
+	);
+	const dependentTypeDeclarations = rawDependentTypeDeclarations.filter(
+		(declaration) => declaration !== platformPropsTypeDeclaration,
 	);
 
 	// 3) extract the import statement
@@ -640,9 +726,20 @@ const generateComponentPropTypeSourceCodeWithSerializedType = (
 	);
 
 	// 5) serialize the prop type for the @atlaskit component (e.g. PlatformButtonProps)
-	const [typeDefCode, usedExternalTypes] = serializeSymbolType(
-		baseComponentPropSymbol,
-		sourceFile.getProject().getTypeChecker(),
+	const [typeDefCode, usedExternalTypes] = serializeTypeReferenceWithPickType(
+		baseComponentPropTypeReference,
+		({ jsDoc, typeCode, propertySignature }) => {
+			const propertyName = propertySignature.getName();
+			if (omitKeys.includes(propertyName)) {
+				return {
+					typeCode,
+				};
+			}
+			return {
+				jsDoc,
+				typeCode,
+			};
+		},
 	);
 
 	// 6) generate the source file
@@ -654,7 +751,7 @@ const generateComponentPropTypeSourceCodeWithSerializedType = (
 	const dependentTypeCode = [
 		...dependentTypeDeclarations.map((typeAlias) => typeAlias.getText()),
 		platformPropsTypeDeclarationName &&
-			`\n// Serialized type\ntype ${platformPropsTypeDeclarationName} = ${typeDefCode}`,
+			`\n// Serialized type\ntype ${platformPropsTypeDeclarationName} = ${typeDefCode};`,
 	]
 		.filter(Boolean)
 		.join('\n');
@@ -725,7 +822,14 @@ const codeConsolidators: Record<string, CodeConsolidator> = {
 	PressableProps: handleXCSSProp,
 };
 
-const typeSerializableComponentPropSymbols = ['CodeProps', 'CodeBlockProps'];
+const typeSerializableComponentPropSymbols = [
+	'CalendarProps',
+	'CodeProps',
+	'CodeBlockProps',
+	'BadgeProps',
+	'HeadingProps',
+	'RangeProps',
+];
 
 const generateComponentPropTypeSourceCode = (
 	componentPropSymbol: Symbol,

package/scripts/codegen/componentPropTypes.ts

@@ -189,17 +189,24 @@ const generateSharedTypesFile = (componentOutputDir: string) => {
 	fs.writeFileSync(typesFilePath, signedSourceCode);
 };
 
-const generateComponentPropTypes = (componentName?: string) => {
+const generateComponentPropTypes = (componentNames?: string) => {
 	const componentOutputDir = resolve(__dirname, '..', '..', 'src', 'components', '__generated__');
 	const componentIndexSourceFile = forgeUIProject.addSourceFileAtPath(
 		require.resolve('@atlassian/forge-ui/UIKit'),
 	);
 	try {
+		const componentNamesFilter = componentNames ? componentNames.split(',') : [];
 		const componentPropTypeSymbols = componentIndexSourceFile
 			.getExportSymbols()
-			.filter((symbol) =>
-				symbol.getName().endsWith(componentName ? `${componentName}Props` : 'Props'),
-			)
+			.filter((symbol) => {
+				if (componentNamesFilter.length === 0) {
+					return symbol.getName().endsWith('Props');
+				}
+				const symbolName = symbol.getName();
+				return componentNamesFilter.some((componentName) => {
+					return symbolName === `${componentName}Props`;
+				});
+			})
 			.sort((a, b) => a.getName().localeCompare(b.getName()));
 
 		// generate share types file first

package/scripts/codegen/typeSerializer.ts

@@ -1,118 +1,137 @@
 import {
-	type Symbol as TSSymbol,
 	type Node,
-	type TypeChecker,
 	type PropertySignature,
 	type Type as TSType,
+	type UnionTypeNode,
+	type TypeReferenceNode,
 	SyntaxKind,
 } from 'ts-morph';
 
-export const serializeSymbolType = (
-	symbol: TSSymbol,
-	typeChecker: TypeChecker,
+type PropertyCallback = ({
+	jsDoc,
+	typeCode,
+	propertySignature,
+}: {
+	jsDoc?: string;
+	typeCode: string;
+	propertySignature: PropertySignature;
+}) => {
+	jsDoc?: string;
+	typeCode: string;
+} | null;
+
+export const serializeTypeReferenceWithPickType = (
+	typeReference: TypeReferenceNode,
+	propertyCallback: PropertyCallback,
 ): [string, Set<string>] => {
-	const declaration = symbol.getDeclarations()[0];
+	const usedExternalTypes = new Set<string>();
+	const serializedType = flattenPickType(typeReference, usedExternalTypes, propertyCallback);
+	return [serializedType, usedExternalTypes];
+};
 
-	if (!declaration) {
-		throw new Error(`No declaration found for symbol: ${symbol.getName()}`);
+// resolve single level type references (e.g. SupportedLanguages))
+// type SupportedLanguages = 'text' | 'PHP' | 'Java' | 'CSharp' | ...;
+const isSimpleTypeReferenceNode = (tsType: TSType): boolean => {
+	if (tsType.isUnion()) {
+		const unionTypes = tsType.getUnionTypes();
+		return unionTypes.every((t) => isBasicType(t));
 	}
+	return false;
+};
 
-	const usedExternalTypes = new Set<string>();
-	if (declaration.getKind() === SyntaxKind.TypeAliasDeclaration) {
-		const typeAlias = declaration.asKindOrThrow(SyntaxKind.TypeAliasDeclaration);
-		const typeNode = typeAlias.getTypeNode();
-		if (typeNode && isCommonComponentPropType(typeNode)) {
-			const serializedType = flattenPickType(
-				typeNode.asKindOrThrow(SyntaxKind.TypeReference),
-				typeChecker,
-				usedExternalTypes,
-			);
-			return [serializedType, usedExternalTypes];
+// event function type that is not a React event handler
+const isCustomEventHandlerType = (tsType: TSType): boolean => {
+	const callSignatures = tsType.getCallSignatures();
+	// we already import React, hence we don't have to serialize it
+	if (
+		callSignatures.length > 0 &&
+		callSignatures[0].getParameters().length > 1 &&
+		!tsType.getText().startsWith('React.')
+	) {
+		const analyticsEventSymbol = callSignatures[0].getParameters()[1];
+		if (analyticsEventSymbol) {
+			const eventTypeName = analyticsEventSymbol
+				.getDeclarations()[0]
+				.asKind(SyntaxKind.Parameter)
+				?.getTypeNode()
+				?.getText();
+			return eventTypeName === 'UIAnalyticsEvent';
 		}
+		return true;
 	}
+	return false;
+};
 
-	throw new Error(
-		`Unsupported declaration kind: ${declaration.getKindName()} for symbol: ${symbol.getName()}`,
-	);
+const serializeSimpleEventType = (tsType: TSType): string => {
+	const propertyCode: string[] = tsType
+		.getProperties()
+		.map((prop) => {
+			const propertySignature = prop.getDeclarations()[0] as PropertySignature | undefined;
+			if (propertySignature) {
+				const typeCode = serializeSimpleTypeNode(resolveNonNullableType(propertySignature));
+				return `${prop.getName()}: ${typeCode}`;
+			}
+			return null;
+		})
+		.filter(Boolean) as string[];
+	return `{ ${propertyCode.join(', ')} }`;
 };
 
-/**
- * Checks if a node is a common component prop type. e.g.
- *
- * export type BleedProps = Pick<
- *   PlatformBleedProps,
- *   'children' | 'all' | 'inline' | 'block' | 'testId' | 'role'
- * >;
- */
-const isCommonComponentPropType = (node: Node) => {
-	if (node.getKind() === SyntaxKind.TypeReference) {
-		const typeRef = node.asKindOrThrow(SyntaxKind.TypeReference);
-		const typeName = typeRef.getTypeName().getText();
-		return typeName === 'Pick';
-	}
-	return false;
+const serializeCustomEventHandlerType = (tsType: TSType): string => {
+	const callSignature = tsType.getCallSignatures()[0];
+	const eventSymbol = callSignature.getParameters()[0];
+	const parameterDeclaration = eventSymbol.getDeclarations()[0].asKindOrThrow(SyntaxKind.Parameter);
+	const eventType = parameterDeclaration.getType();
+	const analyticsEventName = callSignature.getParameters()[1]?.getName() ?? 'analyticsEvent';
+	const returnType = callSignature.getReturnType().getText();
+	return `(${eventSymbol.getName()}: ${serializeSimpleEventType(eventType)}, ${analyticsEventName}: any) => ${returnType}`;
 };
 
-// resolve single level type references (e.g. SupportedLanguages))
-// type SupportedLanguages = 'text' | 'PHP' | 'Java' | 'CSharp' | ...;
-const isSimpleTypeReferenceNode = (node: Node): boolean => {
-	const type = node.getType();
-	if (type.isUnion()) {
-		const unionTypes = type.getUnionTypes();
-		return unionTypes.every((t) => {
-			return t.isString() || t.isStringLiteral() || t.isNumber() || t.isBoolean();
-		});
+const serializeSimpleTypeNode = (tsType: TSType): string => {
+	if (tsType.isStringLiteral()) {
+		return `'${tsType.getLiteralValue()}'`;
+	} else if (isBasicType(tsType)) {
+		return tsType.getText();
+	} else if (isSimpleTypeReferenceNode(tsType)) {
+		const unionTypes = tsType.getUnionTypes();
+		const serializedTypes = unionTypes.map((t) => serializeSimpleTypeNode(t));
+		return serializedTypes.join(' | ');
+	} else if (isCustomEventHandlerType(tsType)) {
+		return serializeCustomEventHandlerType(tsType);
 	}
-	return false;
+	return tsType.getText();
 };
 
-const serializeSimpleTypeNode = (node: Node): string => {
-	switch (node.getKind()) {
-		case SyntaxKind.UnionType:
-			const unionType = node.asKindOrThrow(SyntaxKind.UnionType);
-			const unionTypes = unionType.getTypeNodes().map((t) => serializeSimpleTypeNode(t));
-			return unionTypes.join(' | ');
-
-		case SyntaxKind.StringLiteral:
-			return `'${node.asKindOrThrow(SyntaxKind.StringLiteral).getLiteralValue()}'`;
-		case SyntaxKind.TypeReference:
-			// resolve single level type references (e.g. SupportedLanguages))
-			if (isSimpleTypeReferenceNode(node)) {
-				return node
-					.getType()
-					.getUnionTypes()
-					.map((t) => t.getText())
-					.join(' | ');
-			}
+const resolveNonNullableType = (propertySignature: PropertySignature): TSType => {
+	const hasOptionalHint = propertySignature.hasQuestionToken();
+	const type = propertySignature.getType();
+	if (!hasOptionalHint) {
+		return type;
 	}
-	return node.getText();
+	// there is a case where `children?: ReactNode` using getNonNullableType() don't return the original ReactNode type
+	if (type.getText().split(' | ').includes('undefined')) {
+		return type.getNonNullableType();
+	}
+	return type;
 };
 
 const serializePropertySignatureCode = (propertySignature: PropertySignature) => {
-	return `${propertySignature.getName()}: ${serializeSimpleTypeNode(propertySignature.getTypeNode()!)};`;
+	const propertyName = propertySignature.getName();
+	const isOptional = propertySignature.hasQuestionToken();
+	const typeCode = serializeSimpleTypeNode(resolveNonNullableType(propertySignature));
+	return `${propertyName}${isOptional ? '?' : ''}: ${typeCode};`;
 };
 
 const flattenPickType = (
 	typeRef: Node,
-	typeChecker: TypeChecker,
 	usedExternalTypesOutput: Set<string>,
+	propertyCallback: PropertyCallback,
 ): string => {
-	const typeArgs = typeRef.asKindOrThrow(SyntaxKind.TypeReference).getTypeArguments();
-
-	if (typeArgs.length < 2) {
-		return typeRef.getText(); // Fallback if not a valid Pick
-	}
-
-	const keysNode = typeArgs[1];
-
-	// Get the selected keys
-	const selectedKeys = extractUnionKeys(keysNode);
-
-	// Get the base type symbol
+	const pickKeys = extractPickKeys(typeRef.asKindOrThrow(SyntaxKind.TypeReference));
 	const properties = typeRef
 		.getType()
 		.getProperties()
-		.filter((prop) => selectedKeys.includes(prop.getName()));
+		.filter((prop) => pickKeys.includes(prop.getName()));
 
 	if (properties.length === 0) {
 		return '{}'; // If no properties match, return 'any'
@@ -124,19 +143,23 @@ const flattenPickType = (
 			if (!propertySignature) {
 				return null; // Skip if no declaration
 			}
-			const jsDoc = propertySignature.getJsDocs()?.[0]?.getText();
-			const typeCode = `\t${serializePropertySignatureCode(propertySignature)}`;
+			const { jsDoc, typeCode } =
+				propertyCallback({
+					propertySignature,
+					jsDoc: propertySignature.getJsDocs()?.[0]?.getText(),
+					typeCode: serializePropertySignatureCode(propertySignature),
+				}) || {};
 			getUnresolvableTypes(propertySignature.getType()).forEach((typeName) => {
 				usedExternalTypesOutput.add(typeName);
 			});
-			return [jsDoc, typeCode].filter(Boolean).join('\n');
+			return `${jsDoc ?? ''}\n\t${typeCode}`;
 		})
 		.filter(Boolean);
 
 	if (serializedProperties.length === 0) {
 		return '{}'; // If no properties are serialized, return empty object
 	}
-	return `{\n  ${serializedProperties.join('\n  ')}\n}`;
+	return `{\n${serializedProperties.map((prop) => (!!prop?.trim() ? `  ${prop}` : '')).join('\n')}\n}`;
 };
 
 const getUnresolvableTypes = (tsType: TSType) => {
@@ -161,8 +184,11 @@ const getUnresolvableTypesBase = (tsType: TSType, unresolvableTypes: Set<string>
 const isBasicType = (tsType: TSType): boolean => {
 	return (
 		tsType.isString() ||
+		tsType.isStringLiteral() ||
 		tsType.isNumber() ||
+		tsType.isNumberLiteral() ||
 		tsType.isBoolean() ||
+		tsType.isBooleanLiteral() ||
 		tsType.isNull() ||
 		tsType.isUndefined() ||
 		tsType.isAny() ||
@@ -173,10 +199,14 @@ const isBasicType = (tsType: TSType): boolean => {
 
 const isExternalType = (tsType: TSType): boolean => {
 	const symbol = tsType.getSymbol() ?? tsType.getAliasSymbol();
-	if (!symbol) {return false;}
+	if (!symbol) {
+		return false;
+	}
 
 	const declaration = symbol.getDeclarations()?.[0];
-	if (!declaration) {return false;}
+	if (!declaration) {
+		return false;
+	}
 
 	const sourceFile = declaration.getSourceFile();
 	const filePath = sourceFile.getFilePath();
@@ -184,18 +214,51 @@ const isExternalType = (tsType: TSType): boolean => {
 	return filePath.includes('node_modules');
 };
 
+export const extractUnionKeysFromTypeReferenceNode = (
+	typeReferenceNode: TypeReferenceNode,
+	targetTypeName: 'Pick' | 'Omit',
+): string[] => {
+	const typeName = typeReferenceNode.getTypeName().getText();
+	if (typeName !== targetTypeName) {
+		throw new Error(`Expected '${targetTypeName}' type, but found '${typeName}'`);
+	}
+	const typeArgs = typeReferenceNode.getTypeArguments();
+	if (typeArgs.length !== 2) {
+		throw new Error(
+			`Expected 2 type arguments for ${targetTypeName}, but found ${typeArgs.length}`,
+		);
+	}
+	const unionType = typeArgs[1];
+	return extractUnionKeys(unionType);
+};
+
+export const extractOmitKeys = (typeReferenceNode: TypeReferenceNode): string[] => {
+	return extractUnionKeysFromTypeReferenceNode(typeReferenceNode, 'Omit');
+};
+
+export const extractPickKeys = (typeReferenceNode: TypeReferenceNode): string[] => {
+	return extractUnionKeysFromTypeReferenceNode(typeReferenceNode, 'Pick');
+};
+
+const unwrapStringQuotes = (str: string): string => {
+	return str.replace(/['"]/g, '');
+};
+
+const extractUnionKeysFromUnionType = (unionType: UnionTypeNode): string[] => {
+	return unionType.getTypeNodes().map((node) => {
+		if (node.getKind() === SyntaxKind.StringLiteral) {
+			return node.asKindOrThrow(SyntaxKind.StringLiteral).getLiteralValue();
+		}
+		return unwrapStringQuotes(node.getText());
+	});
+};
+
 const extractUnionKeys = (keysNode: Node): string[] => {
 	if (keysNode.getKind() === SyntaxKind.UnionType) {
 		const unionType = keysNode.asKindOrThrow(SyntaxKind.UnionType);
-		return unionType.getTypeNodes().map((node) => {
-			if (node.getKind() === SyntaxKind.StringLiteral) {
-				return node.asKindOrThrow(SyntaxKind.StringLiteral).getLiteralValue();
-			}
-			return node.getText().replace(/['"]/g, '');
-		});
+		return extractUnionKeysFromUnionType(unionType);
 	} else if (keysNode.getKind() === SyntaxKind.StringLiteral) {
 		return [keysNode.asKindOrThrow(SyntaxKind.StringLiteral).getLiteralValue()];
 	}
-
-	return [keysNode.getText().replace(/['"]/g, '')];
+	return [unwrapStringQuotes(keysNode.getText())];
 };