/*---------------------------------------------------------------------------------------------

* Copyright (c) Microsoft Corporation. All rights reserved.

*--------------------------------------------------------------------------------------------*/

/**

* Python code generator for session-events and RPC types.

*/

import fs from "fs/promises";

import path from "path";

import type { JSONSchema7, JSONSchema7Definition } from "json-schema";

import { fileURLToPath } from "url";

import {

cloneSchemaForCodegen,

filterNodeByVisibility,

fixNullableRequiredRefsInApiSchema,

getApiSchemaPath,

getRpcSchemaTypeName,

getSessionEventsSchemaPath,

isObjectSchema,

isOpaqueJson,

isVoidSchema,

getNullableInner,

isRpcMethod,

isNodeFullyExperimental,

isNodeFullyDeprecated,

isSchemaDeprecated,

isSchemaExperimental,

isSchemaInternal,

postProcessSchema,

propagateInternalVisibility,

collectInternalSymbols,

collectInternalFieldsOnPublicTypes,

annotateInternalPythonFields,

renameInternalPythonSymbols,

stripBooleanLiterals,

writeGeneratedFile,

collectDefinitionCollections,

collectExperimentalOnlyRpcReferencedDefinitionNames,

collectReachableDefinitionNames,

collectRpcMethodReferencedDefinitionNames,

findSharedSchemaDefinitions,

hasSchemaPayload,

parseExternalSchemaRef,

refTypeName,

resolveObjectSchema,

resolveSchema,

rewriteSharedDefinitionReferences,

withSharedDefinitions,

getSessionEventVariantSchemas,

getSharedSessionEventEnvelopeProperties,

getEnumValueDescriptions,

type ApiSchema,

type DefinitionCollections,

type EnumValueDescriptions,

type RpcMethod,

type SessionEventEnvelopeProperty,

} from "./utils.js";

// ── Utilities ───────────────────────────────────────────────────────────────

const EXTERNAL_SCHEMA_PY_MODULE: Record<string, string> = {

"session-events.schema.json": ".session_events",

};

type PyExperimentalSubject = "type" | "enum" | "event";

function pyExperimentalComment(subject: PyExperimentalSubject, indent = ""): string {

return `${indent}# Experimental: this ${subject} is part of an experimental API and may change or be removed.`;

}

function rewriteExternalRefsForPython(schema: JSONSchema7 & { definitions?: Record<string, JSONSchema7> }): {

placeholderNames: Map<string, string>;

imports: Map<string, Set<string>>;

} {

const placeholderNames = new Map<string, string>();

const imports = new Map<string, Set<string>>();

const placeholderFor = (typeName: string): string => `__ExternalRef_${typeName}`;

const visit = (value: unknown): void => {

if (Array.isArray(value)) {

for (const item of value) visit(item);

return;

}

if (!value || typeof value !== "object") return;

const node = value as Record<string, unknown>;

if (typeof node.$ref === "string" && !node.$ref.startsWith("#")) {

const externalRef = parseExternalSchemaRef(node.$ref);

const module = externalRef ? EXTERNAL_SCHEMA_PY_MODULE[externalRef.schemaFile] : undefined;

if (externalRef && module) {

const placeholder = placeholderFor(externalRef.definitionName);

placeholderNames.set(placeholder, externalRef.definitionName);

let bucket = imports.get(module);

if (!bucket) {

bucket = new Set<string>();

imports.set(module, bucket);

}

bucket.add(externalRef.definitionName);

node.$ref = `#/definitions/${placeholder}`;

}

}

for (const child of Object.values(node)) visit(child);

};

visit(schema);

if (placeholderNames.size > 0) {

if (!schema.definitions) schema.definitions = {};

for (const placeholder of placeholderNames.keys()) {

if (!schema.definitions[placeholder]) {

const markerProperty = `__externalRefMarker_${placeholder}`;

schema.definitions[placeholder] = {

type: "object",

additionalProperties: false,

title: placeholder,

properties: {

[markerProperty]: { type: "string" },

},

required: [markerProperty],

};

}

}

}

return { placeholderNames, imports };

}

function placeholderToQuicktypeIdentifier(placeholder: string): string {

return placeholder

.replace(/^_+/, "")

.split("_")

.map((segment) => (segment ? segment[0].toUpperCase() + segment.slice(1) : ""))

.join("");

}

function placeholderToQuicktypeIdentifiers(placeholder: string): string[] {

const basic = placeholderToQuicktypeIdentifier(placeholder);

return [...new Set([basic, basic.replace(/Mcp/g, "MCP")])];

}

function postProcessExternalRefsForPython(

code: string,

placeholderToReal: Map<string, string>,

externalEnumNames: Set<string> = new Set()

): string {

for (const [placeholder, realName] of placeholderToReal) {

for (const quicktypeName of placeholderToQuicktypeIdentifiers(placeholder)) {

code = code.replace(

new RegExp(

`(?:^|\\n)@dataclass\\r?\\nclass ${quicktypeName}\\b[\\s\\S]*?(?=\\n@dataclass\\b|\\nclass\\s+\\w|\\ndef\\s+\\w|$)`,

"g"

),

"\n"

);

code = code.replace(

new RegExp(

`(?:^|\\n)class ${quicktypeName}\\w*\\(Enum\\):[\\s\\S]*?(?=\\nclass\\s+\\w|\\n@dataclass\\b|\\ndef\\s+\\w|$)`,

"g"

),

"\n"

);

code = code.replace(new RegExp(`\\b${quicktypeName}\\b`, "g"), realName);

}

if (externalEnumNames.has(realName)) {

code = code.replace(new RegExp(`\\b${realName}\\.from_dict\\b`, "g"), realName);

code = code.replace(

new RegExp(`to_class\\(${realName},\\s*([^)]+)\\)`, "g"),

`to_enum(${realName}, $1)`

);

}

}

return code.replace(/\n{3,}/g, "\n\n");

}

function collectPythonExternalEnumNames(

schema: JSONSchema7 | undefined,

placeholderToReal: Map<string, string>

): Set<string> {

const enumNames = new Set<string>();

if (!schema) return enumNames;

const definitions = collectDefinitionCollections(schema as Record<string, unknown>);

for (const realName of placeholderToReal.values()) {

const definition = definitions.definitions[realName] ?? definitions.$defs[realName];

const resolved = definition ? resolveSchema(definition, definitions) ?? definition : undefined;

if (

resolved?.enum &&

Array.isArray(resolved.enum) &&

resolved.enum.every((value) => typeof value === "string")

) {

enumNames.add(realName);

}

}

return enumNames;

}

function preservePythonRpcStringDateFields(definitions: Record<string, JSONSchema7>): void {

const quotaSnapshot = definitions.AccountQuotaSnapshot;

const resetDate = quotaSnapshot?.properties?.resetDate as JSONSchema7 | undefined;

if (resetDate?.type === "string" && resetDate.format === "date-time") {

// Keep the existing Python API shape: AccountQuotaSnapshot.reset_date is an ISO string.

delete resetDate.format;

}

}

function collectExternalUnionAliasesForPython(

definitions: Record<string, JSONSchema7>,

placeholderToReal: Map<string, string>

): Map<string, string[]> {

const aliases = new Map<string, string[]>();

for (const [definitionName, definition] of Object.entries(definitions)) {

const variants = definition.anyOf ?? definition.oneOf;

if (!Array.isArray(variants)) continue;

const realNames: string[] = [];

let allExternal = true;

for (const variant of variants) {

if (!variant || typeof variant !== "object") {

allExternal = false;

break;

}

const ref = (variant as JSONSchema7).$ref;

if (!ref?.startsWith("#/definitions/")) {

allExternal = false;

break;

}

const placeholder = ref.slice("#/definitions/".length);

const realName = placeholderToReal.get(placeholder);

if (!realName) {

allExternal = false;

break;

}

realNames.push(realName);

}

if (allExternal && realNames.length > 0) {

aliases.set(definitionName, realNames);

}

}

return aliases;

}

function postProcessExternalUnionAliasesForPython(code: string, aliases: Map<string, string[]>): string {

for (const [aliasName, realNames] of aliases) {

const aliasLine = `${aliasName} = ${realNames.join(" | ")}`;

const classPattern = new RegExp(

`(?:^|\\n)@dataclass\\r?\\nclass ${aliasName}\\b[\\s\\S]*?(?=\\n@dataclass\\b|\\nclass\\s+\\w|\\ndef\\s+\\w|$)`,

"g"

);

if (classPattern.test(code)) {

code = code.replace(classPattern, `\n${aliasLine}\n`);

} else if (!new RegExp(`^${aliasName}\\s*=`, "m").test(code)) {

code = `${aliasLine}\n\n${code}`;

}

code = code.replace(

new RegExp(`${aliasName}\\.from_dict`, "g"),

`(lambda x: from_union([${realNames.map((name) => `${name}.from_dict`).join(", ")}], x))`

);

code = code.replace(

new RegExp(`to_class\\(${aliasName},\\s*([^)]+)\\)`, "g"),

`from_union([${realNames.map((name) => `lambda x: to_class(${name}, x)`).join(", ")}], $1)`

);

}

return code.replace(/\n{3,}/g, "\n\n");

}

/**

* Replace flat-merged dataclasses emitted by quicktype for $ref-based

* discriminated unions with proper Python unions: a `Name = VariantA | ...`

* alias plus a `_load_Name(obj)` dispatcher. Rewrites `Name.from_dict(x)` and

* `to_class(Name, x)` references to use the dispatcher / per-variant

* `.to_dict()` so callers transparently get the proper union shape.

*

* Detection: walk top-level definitions and pick those whose schema is an

* `anyOf`/`oneOf` of `$ref`s with a shared `const` discriminator. For each

* such definition we expect quicktype to already have emitted both the

* merged blob (which we'll delete) and the per-variant classes (which we

* keep).

*

* Returns the rewritten types-section code and the list of resolved union

* names; callers can re-apply `applyUnionRewritesToPython` to subsequently

* generated code (e.g. RPC method wrappers) so `Name.from_dict(x)` calls

* there also route through the new dispatcher.

*/

interface ResolvedRefBasedUnion {

aliasName: string;

discriminatorProp: string;

dispatch: Array<{ value: string; typeName: string }>;

}

function postProcessRefBasedDiscriminatedUnionsForPython(

code: string,

definitions: Record<string, JSONSchema7>,

definitionCollections: DefinitionCollections

): { code: string; unions: ResolvedRefBasedUnion[] } {

interface UnionInfo {

aliasName: string;

variantNames: string[];

discriminatorProp: string;

dispatch: Array<{ value: string; typeName: string }>;

description: string | undefined;

}

const unions: UnionInfo[] = [];

for (const [defName, definition] of Object.entries(definitions)) {

const variants = (definition.anyOf ?? definition.oneOf) as JSONSchema7[] | undefined;

if (!Array.isArray(variants) || variants.length < 2) continue;

if (!variants.every((v) => typeof v === "object" && v !== null && typeof v.$ref === "string")) {

continue;

}

const variantRefNames = variants.map((v) => refTypeName(v.$ref as string, definitionCollections));

const resolvedVariants = variants.map(

(v) =>

resolveObjectSchema(v, definitionCollections) ??

resolveSchema(v, definitionCollections) ??

v

);

if (resolvedVariants.some((rv) => !rv || rv.properties === undefined)) continue;

const discriminator = findPyDiscriminator(resolvedVariants as JSONSchema7[]);

if (!discriminator) continue;

const aliasName = toPascalCase(defName);

const dispatch = variants.map((_, i) => {

const discProp = (resolvedVariants[i].properties as Record<string, JSONSchema7>)[

discriminator.property

];

return {

value: String(discProp.const),

typeName: toPascalCase(variantRefNames[i]),

};

});

unions.push({

aliasName,

variantNames: variantRefNames.map(toPascalCase),

discriminatorProp: discriminator.property,

dispatch,

description: typeof definition.description === "string" ? definition.description : undefined,

});

}

const resolved: ResolvedRefBasedUnion[] = [];

if (unions.length === 0) return { code, unions: resolved };

const emittedClassNames = new Set<string>();

for (const match of code.matchAll(/^class (\w+)[:\(]/gm)) {

emittedClassNames.add(match[1]);

}

const acronymCandidates = (name: string): string[] => {

const substitutions: Array<[RegExp, string]> = [

[/Api/g, "API"],

[/Mcp/g, "MCP"],

[/Url/g, "URL"],

[/Json/g, "JSON"],

[/Http/g, "HTTP"],

[/Hmac/g, "HMAC"],

[/Tcp/g, "TCP"],

[/Sql/g, "SQL"],

[/Id\b/g, "ID"],

[/Llm/g, "LLM"],

[/Cli/g, "CLI"],

];

const results = new Set<string>([name]);

for (const [pattern, replacement] of substitutions) {

for (const existing of [...results]) {

results.add(existing.replace(pattern, replacement));

}

}

return [...results];

};

const resolveActualName = (expected: string): string | undefined => {

for (const candidate of acronymCandidates(expected)) {

if (emittedClassNames.has(candidate)) return candidate;

}

return undefined;

};

for (const union of unions) {

const actualAliasName = resolveActualName(union.aliasName);

const actualVariantNames: string[] = [];

const actualDispatch: Array<{ value: string; typeName: string }> = [];

let allResolved = true;

for (let i = 0; i < union.variantNames.length; i++) {

const actual = resolveActualName(union.variantNames[i]);

if (!actual) {

allResolved = false;

break;

}

actualVariantNames.push(actual);

actualDispatch.push({ value: union.dispatch[i].value, typeName: actual });

}

if (!allResolved || !actualAliasName) {

continue;

}

resolved.push({

aliasName: actualAliasName,

discriminatorProp: union.discriminatorProp,

dispatch: actualDispatch,

});

const lines = code.split("\n");

let classStart = -1;

for (let i = 0; i < lines.length; i++) {

if (lines[i] === `class ${actualAliasName}:` || lines[i].startsWith(`class ${actualAliasName}(`)) {

classStart = i;

break;

}

}

if (classStart >= 0) {

let blockStart = classStart;

while (

blockStart > 0 &&

(lines[blockStart - 1] === "@dataclass" || /^# /.test(lines[blockStart - 1]))

) {

blockStart--;

}

let blockEnd = classStart + 1;

while (blockEnd < lines.length) {

const ln = lines[blockEnd];

if (

/^class \w/.test(ln) ||

/^def \w/.test(ln) ||

ln === "@dataclass" ||

/^# (?:Experimental|Deprecated|Internal):/.test(ln)

) {

break;

}

blockEnd++;

}

lines.splice(blockStart, blockEnd - blockStart);

code = lines.join("\n");

}

const aliasLine = union.description

? `# ${union.description.replace(/\n/g, " ")}\n${actualAliasName} = ${actualVariantNames.join(" | ")}`

: `${actualAliasName} = ${actualVariantNames.join(" | ")}`;

const dispatcherLines: string[] = [];

dispatcherLines.push(`def _load_${actualAliasName}(obj: Any) -> "${actualAliasName}":`);

dispatcherLines.push(` assert isinstance(obj, dict)`);

dispatcherLines.push(` kind = obj.get(${JSON.stringify(union.discriminatorProp)})`);

dispatcherLines.push(` match kind:`);

for (const m of actualDispatch) {

dispatcherLines.push(` case ${JSON.stringify(m.value)}: return ${m.typeName}.from_dict(obj)`);

}

dispatcherLines.push(

` case _: raise ValueError(f"Unknown ${actualAliasName} ${union.discriminatorProp}: {kind!r}")`

);

code = `${code.trimEnd()}\n\n\n${aliasLine}\n\n\n${dispatcherLines.join("\n")}\n`;

}

code = applyUnionRewritesToPython(code, resolved);

return { code, unions: resolved };

}

/**

* Rewrite occurrences of `Name.from_dict(...)` to `_load_Name(...)` and

* `to_class(Name, x)` to `(x).to_dict()` for each union the caller passes in.

* Safe to apply repeatedly — re-running on already-rewritten code is a no-op.

*/

function applyUnionRewritesToPython(code: string, unions: ResolvedRefBasedUnion[]): string {

for (const union of unions) {

code = code.replace(

new RegExp(`\\b${union.aliasName}\\.from_dict\\b`, "g"),

`_load_${union.aliasName}`

);

code = code.replace(

new RegExp(`to_class\\(${union.aliasName},\\s*([^,)]+)\\)`, "g"),

`($1).to_dict()`

);

}

return code;

}

/**

* For each discriminated-union variant class, replace the dataclass-level

* discriminator field (e.g. ``kind: PermissionDecisionApproveOnceKind``) with

* a class-level constant (e.g. ``kind: ClassVar[str] = "approve-once"``).

* This lets users construct variants without supplying the discriminator

* value (``PermissionDecisionApproveOnce()`` instead of

* ``PermissionDecisionApproveOnce(kind=PermissionDecisionApproveOnceKind.APPROVE_ONCE)``),

* matching the TS / Rust / .NET / Go ergonomics for the same schema.

*

* Also rewrites the generated ``from_dict`` to skip parsing the discriminator

* (the dispatcher routed based on it; the variant class identity carries it)

* and ``to_dict`` to emit the constant directly.

*/

function postProcessDiscriminatorDefaultsForPython(

code: string,

unions: ResolvedRefBasedUnion[]

): string {

// Build variant lookup: variant class name → { prop, value }.

const variantInfo = new Map<string, { prop: string; value: string }>();

for (const union of unions) {

for (const d of union.dispatch) {

// First-wins; multiple unions referencing the same variant share a

// discriminator/value pair anyway.

if (!variantInfo.has(d.typeName)) {

variantInfo.set(d.typeName, { prop: union.discriminatorProp, value: d.value });

}

}

}

if (variantInfo.size === 0) return code;

const lines = code.split("\n");

const out: string[] = [];

let usedClassVar = false;

let i = 0;

while (i < lines.length) {

const line = lines[i];

const classMatch = line.match(/^class (\w+)[:\(]/);

if (!classMatch) {

out.push(line);

i++;

continue;

}

const className = classMatch[1];

const info = variantInfo.get(className);

if (!info) {

out.push(line);

i++;

continue;

}

// Find the bounds of this class block: everything indented under it.

const classStart = i;

let classEnd = i + 1;

while (classEnd < lines.length) {

const ln = lines[classEnd];

if (

/^class \w/.test(ln) ||

/^def \w/.test(ln) ||

ln === "@dataclass" ||

/^# (?:Experimental|Deprecated|Internal):/.test(ln) ||

ln.startsWith("@dataclass(")

) {

break;

}

classEnd++;

}

const block = lines.slice(classStart, classEnd);

// Locate the discriminator field declaration. Quicktype emits

// ` kind: PermissionDecisionApproveOnceKind` while the

// session-events codegen emits ` kind: str` — both match the

// simple `<indent><prop>: <Type>` shape (no default value, since the

// field is required in the schema).

const fieldPattern = new RegExp(`^(\\s+)${info.prop}: [\\w\\[\\], ]+$`);

let fieldIdx = -1;

for (let j = 1; j < block.length; j++) {

if (fieldPattern.test(block[j])) {

fieldIdx = j;

break;

}

}

if (fieldIdx < 0) {

// Variant class without an explicit discriminator field — leave alone.

out.push(...block);

i = classEnd;

continue;

}

const fieldIndent = (block[fieldIdx].match(/^(\s+)/) ?? ["", ""])[1];

const literal = JSON.stringify(info.value);

// Replace the field with a class-level constant.

block[fieldIdx] = `${fieldIndent}${info.prop}: ClassVar[str] = ${literal}`;

usedClassVar = true;

// Drop any field-trailing docstring lines that immediately followed the

// original field. Quicktype emits """..."""-style block strings; the

// session-events codegen does not emit per-field docstrings. We only

// touch the line at fieldIdx+1 if it's a docstring or blank.

// (Conservative: leave additional lines in place; they don't reference

// the dropped enum.)

// Rewrite from_dict / to_dict bodies.

for (let j = fieldIdx + 1; j < block.length; j++) {

const ln = block[j];

// Drop `<prop> = ...(obj.get("<prop>"))` parse line in from_dict.

const propAssignPattern = new RegExp(

`^\\s+${info.prop} = .+\\(obj\\.get\\(${JSON.stringify(info.prop)}\\)\\)`

);

if (propAssignPattern.test(ln)) {

block[j] = "<<<DROP>>>";

continue;

}

// Drop multi-line constructor kwarg of the form ` kind=kind,` —

// emitted by the session-events codegen when the constructor call

// is broken across lines.

const multilineKwargPattern = new RegExp(

`^\\s+${info.prop}=${info.prop},?\\s*$`

);

if (multilineKwargPattern.test(ln)) {

block[j] = "<<<DROP>>>";

continue;

}

// Convert `return X(a, prop, b)` (single-line positional) to drop

// the prop arg. Quicktype-emitted constructors are single-line.

const ctorMatch = ln.match(new RegExp(`^(\\s+)return ${className}\\((.*)\\)\\s*$`));

if (ctorMatch) {

const argList = ctorMatch[2];

const args = splitTopLevelCommasMulti(argList);

const filtered = args

.map((a) => a.trim())

.filter((a) => {

const kw = a.match(/^([a-zA-Z_]\w*)\s*=/);

const name = kw ? kw[1] : a;

return name !== info.prop;

});

block[j] = `${ctorMatch[1]}return ${className}(${filtered.join(", ")})`;

continue;

}

// Rewrite `result["<prop>"] = to_enum(<TypeName>, self.<prop>)` to

// emit the class-level constant directly.

const toDictPattern = new RegExp(

`^(\\s+)result\\[${JSON.stringify(info.prop)}\\] = .+`

);

if (toDictPattern.test(ln)) {

const indent = (ln.match(/^(\s+)/) ?? ["", ""])[1];

block[j] = `${indent}result[${JSON.stringify(info.prop)}] = self.${info.prop}`;

continue;

}

}

out.push(...block.filter((l) => l !== "<<<DROP>>>"));

i = classEnd;

}

let result = out.join("\n");

if (usedClassVar) {

result = ensureClassVarImport(result);

}

return result;

}

function splitTopLevelCommasMulti(s: string): string[] {

const parts: string[] = [];

let depth = 0;

let start = 0;

for (let i = 0; i < s.length; i++) {

const c = s[i];

if (c === "(" || c === "[" || c === "{") depth++;

else if (c === ")" || c === "]" || c === "}") depth--;

else if (c === "," && depth === 0) {

parts.push(s.slice(start, i));

start = i + 1;

}

}

parts.push(s.slice(start));

return parts.filter((p) => p.trim().length > 0);

}

function ensureClassVarImport(code: string): string {

// Already imported?

if (/\bfrom typing import [^\n]*\bClassVar\b/.test(code)) return code;

return code.replace(

/^from typing import (.+)$/m,

(_match, names) => {

const list = names.split(",").map((n: string) => n.trim()).filter(Boolean);

list.push("ClassVar");

list.sort();

return `from typing import ${[...new Set(list)].join(", ")}`;

}

);

}

function pushPyExperimentalComment(lines: string[], subject: PyExperimentalSubject, indent = ""): void {

lines.push(pyExperimentalComment(subject, indent));

}

function pushPyExperimentalApiGroupComment(lines: string[]): void {

lines.push("# Experimental: this API group is experimental and may change or be removed.");

}

/**

* Emit `# Deprecated:` / `# Experimental:` / `# Internal:` comments above a

* dataclass field. Order matches our other codegens (deprecated, experimental,

* internal) and keeps the comments out of the field declaration itself.

*/

function pushPyFieldMarkers(lines: string[], propSchema: JSONSchema7 | null | undefined): void {

if (!propSchema) return;

if (isSchemaDeprecated(propSchema)) {

lines.push(` # Deprecated: this field is deprecated.`);

}

if (isSchemaExperimental(propSchema)) {

lines.push(` # Experimental: this field is part of an experimental API and may change or be removed.`);

}

if (isSchemaInternal(propSchema)) {

lines.push(` # Internal: this field is an internal SDK API and is not part of the public surface.`);

}

}

/**

* Modernize quicktype's Python 3.7 output to Python 3.11+ syntax:

* - Optional[T] → T | None

* - List[T] → list[T]

* - Dict[K, V] → dict[K, V]

* - Type[T] → type[T]

* - Callable from collections.abc instead of typing

* - Clean up unused typing imports

*/

function replaceBalancedBrackets(code: string, prefix: string, replacer: (inner: string) => string): string {

let result = "";

let i = 0;

while (i < code.length) {

const idx = code.indexOf(prefix + "[", i);

if (idx === -1) {

result += code.slice(i);

break;

}

result += code.slice(i, idx);

const start = idx + prefix.length + 1; // after '['

let depth = 1;

let j = start;

while (j < code.length && depth > 0) {

if (code[j] === "[") depth++;

else if (code[j] === "]") depth--;

j++;

}

const inner = code.slice(start, j - 1);

result += replacer(inner);

i = j;

}

return result;

}

/** Split a string by commas, but only at the top bracket depth (ignores commas inside [...]) */

function splitTopLevelCommas(s: string): string[] {

const parts: string[] = [];

let depth = 0;

let start = 0;

for (let i = 0; i < s.length; i++) {

if (s[i] === "[") depth++;

else if (s[i] === "]") depth--;

else if (s[i] === "," && depth === 0) {

parts.push(s.slice(start, i));

start = i + 1;

}

}

parts.push(s.slice(start));

return parts;

}

function pyDocstringLiteral(text: string): string {

const normalized = text

.split(/\r?\n/)

.map((line) => line.replace(/\s+$/g, ""))

.join("\n");

return JSON.stringify(normalized);

}

function rpcResultDescription(method: RpcMethod, resultSchema: JSONSchema7 | undefined): string | undefined {

if (isVoidSchema(resultSchema)) return undefined;

return method.result?.description ?? resultSchema?.description;

}

function rpcParamsDescription(method: RpcMethod, effectiveParams: JSONSchema7 | undefined): string | undefined {

return method.params?.description ?? effectiveParams?.description;

}

function pushPyRpcMethodDocstring(

lines: string[],

indent: string,

method: RpcMethod,

options: {

paramsName?: string;

paramsDescription?: string;

resultDescription?: string;

deprecated?: boolean;

experimental?: boolean;

internal?: boolean;

} = {}

): void {

const sections: string[] = [method.description ?? `Calls ${method.rpcMethod}.`];

if (options.paramsName && options.paramsDescription) {

sections.push(`Args:\n ${options.paramsName}: ${options.paramsDescription}`);

}

if (options.resultDescription) {

sections.push(`Returns:\n ${options.resultDescription}`);

}

if (options.deprecated) {

sections.push(".. deprecated:: This API is deprecated and will be removed in a future version.");

}

if (options.experimental) {

sections.push(".. warning:: This API is experimental and may change or be removed in future versions.");

}

if (options.internal) {

sections.push(":meta private:\n\nInternal SDK API; not part of the public surface.");

}

lines.push(`${indent}${pyDocstringLiteral(sections.join("\n\n"))}`);

}

function modernizePython(code: string): string {

// Replace Optional[X] with X | None (handles arbitrarily nested brackets)

code = replaceBalancedBrackets(code, "Optional", (inner) => `${inner} | None`);

// Replace Union[X, Y] with X | Y (split only at top-level commas, not inside brackets)

// Run iteratively to handle nested Union inside Dict/List

let prev = "";

while (prev !== code) {

prev = code;

code = replaceBalancedBrackets(code, "Union", (inner) => {

return splitTopLevelCommas(inner).map((s: string) => s.trim()).join(" | ");

});

}

// Replace List[X] with list[X]

code = code.replace(/\bList\[/g, "list[");

// Replace Dict[K, V] with dict[K, V]

code = code.replace(/\bDict\[/g, "dict[");

// Replace Type[T] with type[T]

code = code.replace(/\bType\[/g, "type[");

// Move Callable from typing to collections.abc

code = code.replace(

/from typing import (.*), Callable$/m,

"from typing import $1\nfrom collections.abc import Callable"

);

code = code.replace(

/from typing import Callable, (.*)$/m,

"from typing import $1\nfrom collections.abc import Callable"

);

// Remove now-unused imports from typing (Optional, List, Dict, Type)

code = code.replace(/from typing import (.+)$/m, (_match, imports: string) => {

const items = imports.split(",").map((s: string) => s.trim());

const remove = new Set(["Optional", "List", "Dict", "Type", "Union"]);

const kept = items.filter((i: string) => !remove.has(i));

return `from typing import ${kept.join(", ")}`;

});

return code;

}

/**

* Collapse lambdas that only forward their single argument into another callable.

* This keeps the generated Python readable and avoids CodeQL "unnecessary lambda" findings.

*/

function unwrapRedundantPythonLambdas(code: string): string {

return code.replace(

/lambda\s+([A-Za-z_][A-Za-z0-9_]*)\s*:\s*((?:[A-Za-z_][A-Za-z0-9_]*)(?:\.[A-Za-z_][A-Za-z0-9_]*)*)\(\1\)/g,

"$2"

);

}

function collapsePlaceholderPythonDataclasses(code: string, knownDefinitionNames?: Set<string>): string {

const classBlockRe = /(@dataclass\r?\nclass\s+(\w+):[\s\S]*?)(?=^@dataclass|^class\s+\w+|^def\s+\w+|\Z)/gm;

const matches = [...code.matchAll(classBlockRe)].map((match) => ({

fullBlock: match[1],

name: match[2],

normalizedBody: normalizePythonDataclassBlock(match[1], match[2]),

}));

const groups = new Map<string, typeof matches>();

for (const match of matches) {

const group = groups.get(match.normalizedBody) ?? [];

group.push(match);

groups.set(match.normalizedBody, group);

}

for (const group of groups.values()) {

if (group.length < 2) continue;

const canonical = chooseCanonicalPlaceholderDuplicate(group.map(({ name }) => name), knownDefinitionNames);

if (!canonical) continue;

for (const duplicate of group) {

if (duplicate.name === canonical) continue;

// Only collapse types that quicktype invented (Class suffix or not

// in the schema's named definitions). Preserve intentionally-named types.

if (!isPlaceholderTypeName(duplicate.name) && knownDefinitionNames?.has(duplicate.name.toLowerCase())) continue;

code = code.replace(duplicate.fullBlock, "");

code = code.replace(new RegExp(`\\b${duplicate.name}\\b`, "g"), canonical);

}

}

return code.replace(/\n{3,}/g, "\n\n");

}

/**

* Reorder Python class/enum definitions so forward references are resolved.

* Quicktype may emit classes in an order where a class references another

* that hasn't been defined yet, causing NameError at import time.

* This performs a topological sort of type definitions while preserving

* the relative position of non-class blocks (functions, standalone code).

*/

function reorderPythonForwardRefs(code: string): string {

// Split code into top-level blocks. Each block starts at an unindented

// line that begins a class, decorated class, enum, or function definition.

const lines = code.split("\n");

interface Block {

name: string;

code: string;

isType: boolean; // true for class/enum definitions

}

const blocks: Block[] = [];

let currentLines: string[] = [];

let currentName: string | null = null;

let isType = false;

function flushBlock() {

if (currentLines.length === 0) return;

const blockCode = currentLines.join("\n");

blocks.push({

name: currentName ?? `__anon_${blocks.length}`,

code: blockCode,

isType,

});

currentLines = [];

currentName = null;

isType = false;

}

for (let i = 0; i < lines.length; i++) {

const line = lines[i];

const isTopLevel = line.length > 0 && line[0] !== " " && line[0] !== "\t";

if (isTopLevel) {

const classMatch = line.match(/^class\s+(\w+)/);

const defMatch = line.match(/^def\s+(\w+)/);

const decoratorMatch = line === "@dataclass";

const commentMatch = line.startsWith("# ");

if (classMatch) {

// If previous block was just a decorator waiting for a class, merge

if (currentLines.length > 0 && currentName === null && isType) {

// This is the class line following @dataclass

currentName = classMatch[1];

currentLines.push(line);

continue;

}

flushBlock();

currentLines = [line];

currentName = classMatch[1];

isType = true;

} else if (decoratorMatch) {

flushBlock();

currentLines = [line];

isType = true;

} else if (defMatch) {

flushBlock();

currentLines = [line];

currentName = defMatch[1];

isType = false;

} else if (commentMatch && currentLines.length === 0) {

// Standalone comment — attach to next block

currentLines = [line];

} else {

currentLines.push(line);

}

} else {

currentLines.push(line);

}

}

flushBlock();

if (blocks.length === 0) return code;

// Collect all type names (classes and enums)

const typeNames = new Set(blocks.filter((b) => b.isType).map((b) => b.name));

if (typeNames.size === 0) return code;

// Build dependency graph: for each type block, find references to other type names

const deps = new Map<string, Set<string>>();

for (const block of blocks) {

if (!block.isType) continue;

const blockDeps = new Set<string>();

for (const tn of typeNames) {

if (tn === block.name) continue;

if (new RegExp(`\\b${tn}\\b`).test(block.code)) {

blockDeps.add(tn);

}

}

deps.set(block.name, blockDeps);

}

// Kahn's algorithm for topological sort

const inDegree = new Map<string, number>();

for (const tn of typeNames) inDegree.set(tn, deps.get(tn)?.size ?? 0);

const dependents = new Map<string, string[]>();

for (const tn of typeNames) dependents.set(tn, []);