A codemod that adds TypedDocumentNode<any, any> typing to Apollo query and mutate calls

transform-cast-client-query-to-tdn.js

// Robust patcher for Apollo Client calls:
// - Edits client.query({...}) and client.mutate({...})
// - If options arg is an ObjectExpression, casts its `query:` or `mutation:` value to
//     `TypedDocumentNode<any, any>` and inserts a type-only import.
// - If options arg is an Identifier bound to an object literal in the same file, edits that.
// - Otherwise, falls back to adding generics to the call: <any, any>
// Works with awaited and non-awaited calls. Also tolerates optional chaining.
// Run with: --parser=tsx --extensions=ts,tsx
module.exports.parser = "tsx";

module.exports = function transformer(file, api) {
	const j = api.jscodeshift;
	const root = j(file.source);

	const TDN_SOURCE = "@graphql-typed-document-node/core";
	const TDN_NAME = "TypedDocumentNode";
	let didInsertTDNImport = false;
	let didAnyChange = false;

	// ---- helpers ----
	const makeTDNAnyAny = () =>
		j.tsTypeReference(
			j.identifier(TDN_NAME),
			j.tsTypeParameterInstantiation([j.tsAnyKeyword(), j.tsAnyKeyword()]),
		);

	const castToTDN = (expr) => {
		if (expr.type === "TSAsExpression" || expr.type === "TypeAssertion")
			return expr;
		didInsertTDNImport = true;
		return j.tsAsExpression(expr, makeTDNAnyAny());
	};

	const isClientCall = (node) => {
		// Match .query(...) or .mutate(...)
		if (!node || node.type !== "CallExpression") return null;

		const callee = node.callee;
		// Identifier (rare): query(...)
		if (
			callee.type === "Identifier" &&
			(callee.name === "query" || callee.name === "mutate")
		) {
			return callee.name;
		}
		// MemberExpression: client.query(...), Apollo.mutate(...)
		if (
			callee.type === "MemberExpression" &&
			callee.property.type === "Identifier"
		) {
			const name = callee.property.name;
			if (name === "query" || name === "mutate") return name;
		}
		// OptionalMemberExpression (if using optional chaining)
		if (
			callee.type === "OptionalMemberExpression" &&
			callee.property.type === "Identifier"
		) {
			const name = callee.property.name;
			if (name === "query" || name === "mutate") return name;
		}
		return null;
	};

	const hasTypeArgsOrCastedFirstArg = (call) => {
		if (call.typeArguments) return true;
		const firstArg = call.arguments && call.arguments[0];
		if (!firstArg) return false;
		return (
			firstArg.type === "TSAsExpression" || firstArg.type === "TypeAssertion"
		);
	};

	const findObjectProp = (objExpr, keys) => {
		// keys: ['query'] or ['mutation']
		return (objExpr.properties || []).find((p) => {
			if (!p || p.type !== "Property" || p.computed) return false;
			// identifier key
			if (p.key.type === "Identifier" && keys.includes(p.key.name)) return true;
			// string literal key: 'query'
			if (
				p.key.type === "Literal" &&
				typeof p.key.value === "string" &&
				keys.includes(p.key.value)
			)
				return true;
			// TS literal: "query"
			if (p.key.type === "StringLiteral" && keys.includes(p.key.value))
				return true;
			return false;
		});
	};

	// Try to resolve an Identifier to a variable declarator with an ObjectExpression initializer
	const resolveIdentifierObjectInit = (name) => {
		let found = null;
		root
			.find(j.VariableDeclarator, {
				id: { type: "Identifier", name },
				init: { type: "ObjectExpression" },
			})
			.forEach((p) => {
				if (!found) found = p; // pick the first
			});
		return found;
	};

	// ---- main pass: edit call expressions ----
	root.find(j.CallExpression).forEach((path) => {
		const node = path.value;
		const kind = isClientCall(node); // 'query' | 'mutate' | null
		if (!kind) return;

		// If generics or cast already applied on the call, skip
		if (hasTypeArgsOrCastedFirstArg(node)) return;

		// options arg
		const firstArg = node.arguments && node.arguments[0];
		const key = kind === "query" ? "query" : "mutation";

		// Case A: Inline object literal
		if (firstArg && firstArg.type === "ObjectExpression") {
			const prop = findObjectProp(firstArg, [key]);
			if (prop && prop.value) {
				prop.value = castToTDN(prop.value);
				didAnyChange = true;
				return;
			}
			// If no prop (odd), fall back to generics
		}

		// Case B: Identifier referencing an object literal variable
		if (firstArg && firstArg.type === "Identifier") {
			const decl = resolveIdentifierObjectInit(firstArg.name);
			if (decl) {
				const objExpr = decl.value.init;
				const prop = findObjectProp(objExpr, [key]);
				if (prop && prop.value) {
					prop.value = castToTDN(prop.value);
					didAnyChange = true;
					return;
				}
				// No prop found -> fall back to generics below
			}
		}

		// Case C: Fallback — add generics to the call itself
		const anyT = j.tsAnyKeyword();
		node.typeArguments = j.tsTypeParameterInstantiation([anyT, anyT]);
		didAnyChange = true;
	});

	// Insert type-only import if we added any casts
	if (didAnyChange && didInsertTDNImport) {
		const existing = root.find(j.ImportDeclaration, {
			source: { value: TDN_SOURCE },
		});
		if (existing.size() > 0) {
			existing.forEach((p) => {
				p.value.importKind = "type";
				const specs = p.value.specifiers || (p.value.specifiers = []);
				const hasTDN = specs.some(
					(s) =>
						s.type === "ImportSpecifier" &&
						s.imported &&
						s.imported.name === TDN_NAME,
				);
				if (!hasTDN) {
					specs.push(j.importSpecifier(j.identifier(TDN_NAME)));
				}
			});
		} else {
			const decl = j.importDeclaration(
				[j.importSpecifier(j.identifier(TDN_NAME))],
				j.literal(TDN_SOURCE),
			);
			decl.importKind = "type";
			const firstImport = root.find(j.ImportDeclaration).at(0);
			if (firstImport.size()) firstImport.insertBefore(decl);
			else root.get().node.program.body.unshift(decl);
		}
	}

	return didAnyChange
		? root.toSource({ quote: "single", trailingComma: true })
		: file.source;
};