@graphql-codegen/python 0.0.0-stage → 0.0.1

package/src/visitor.ts

@@ -0,0 +1,346 @@
+import autoBind from 'auto-bind';
+import {
+  EnumTypeDefinitionNode,
+  FieldDefinitionNode,
+  GraphQLSchema,
+  InputObjectTypeDefinitionNode,
+  InputValueDefinitionNode,
+  InterfaceTypeDefinitionNode,
+  ListTypeNode,
+  NamedTypeNode,
+  NonNullTypeNode,
+  ObjectTypeDefinitionNode,
+  UnionTypeDefinitionNode,
+} from 'graphql';
+import { isValidClassAttr, PYTHON_SCALARS, pythonFieldName } from '@graphql-codegen/python-common';
+import {
+  BaseTypesVisitor,
+  getConfigValue,
+  ParsedTypesConfig,
+} from '@graphql-codegen/visitor-plugin-common';
+import { PythonPluginConfig } from './config.js';
+
+export { PYTHON_SCALARS };
+
+const INDENT = '    ';
+
+export type PythonNamingConvention = 'keep' | 'snakeCase';
+
+export interface PythonPluginParsedConfig extends ParsedTypesConfig {
+  baseModel: string;
+  runtimeImport: string;
+  pyNamingConvention: PythonNamingConvention;
+}
+
+export class PythonVisitor<
+  TRawConfig extends PythonPluginConfig = PythonPluginConfig,
+  TParsedConfig extends PythonPluginParsedConfig = PythonPluginParsedConfig,
+> extends BaseTypesVisitor<TRawConfig, TParsedConfig> {
+  private _usesEnum = false;
+  private _usesTypedDict = false;
+  private _usesBaseModel = false;
+  private _usesUnion = false;
+
+  constructor(schema: GraphQLSchema, pluginConfig: TRawConfig) {
+    // `namingConvention` here is our own field-naming option (`'keep'` |
+    // `'snakeCase'`) and must NOT leak into the base visitor's `namingConvention`
+    // (which controls TYPE-name conversion via `convertFactory`). Strip it before
+    // calling `super` so type names keep their default (PascalCase) conversion.
+    super(
+      schema,
+      // Route custom scalars (no explicit mapping) to `Any` instead of `unknown`.
+      {
+        ...pluginConfig,
+        namingConvention: undefined,
+        defaultScalarType: getConfigValue(pluginConfig.defaultScalarType, 'Any'),
+      },
+      {
+        baseModel: getConfigValue(pluginConfig.baseModel, 'BaseModel'),
+        runtimeImport: getConfigValue(pluginConfig.runtimeImport, 'pydantic'),
+        pyNamingConvention: getConfigValue(pluginConfig.namingConvention, 'snakeCase'),
+      } as TParsedConfig,
+      PYTHON_SCALARS,
+    );
+
+    autoBind(this);
+  }
+
+  /** Python has no `export` keyword; disable the C-family export prefix. */
+  protected getExportPrefix(): string {
+    return '';
+  }
+
+  protected _getScalar(name: string, type: 'input' | 'output'): string {
+    const mapped = this.scalars[name];
+    if (mapped) {
+      return mapped[type];
+    }
+    return 'Any';
+  }
+
+  // --- Nullability / list composition (emit Python instead of TS) ---
+
+  NamedType(node: NamedTypeNode, key: any, parent: any, path: any, ancestors: any): string {
+    return `Optional[${super.NamedType(node, key, parent, path, ancestors)}]`;
+  }
+
+  ListType(node: ListTypeNode, key: any, parent: any, path: any, ancestors: any): string {
+    return `Optional[${super.ListType(node, key, parent, path, ancestors)}]`;
+  }
+
+  protected wrapWithListType(str: string): string {
+    return `list[${str}]`;
+  }
+
+  NonNullType(node: NonNullTypeNode): string {
+    const baseValue = super.NonNullType(node);
+    if (baseValue.startsWith('Optional[') && baseValue.endsWith(']')) {
+      return baseValue.slice('Optional['.length, -1);
+    }
+    return baseValue;
+  }
+
+  // --- Field rendering (name: 'annotation') ---
+
+  private snakeFieldName(name: string): string {
+    return pythonFieldName(name);
+  }
+
+  /**
+   * Field key for a `TypedDict` (input type). Input types are what the caller
+   * SENDS — the runtime forwards the dict to the server as-is — so they ALWAYS
+   * keep the wire (GraphQL) name regardless of `namingConvention`. If the name
+   * is a keyword / non-identifier the containing `TypedDict` is later emitted
+   * with functional syntax (see `buildTypedDict`).
+   */
+  private dictFieldName(name: string): string {
+    return name;
+  }
+
+  /**
+   * Field key for an object / interface type (a class attribute). Under `'keep'`
+   * a keyword / non-identifier name cannot be a class attribute and BaseModel
+   * cannot use functional `TypedDict` syntax, so we fall back to appending `_`
+   * for just that field.
+   */
+  private objectFieldName(name: string): string {
+    if (this.config.pyNamingConvention === 'snakeCase') {
+      return this.snakeFieldName(name);
+    }
+    return isValidClassAttr(name) ? name : `${name}_`;
+  }
+
+  FieldDefinition(node: FieldDefinitionNode): string {
+    if (this.config.onlyEnums) return '';
+    const wire = node.name.value;
+    const attr = this.objectFieldName(wire);
+    // The rendered annotation (e.g. `Optional[str]`, `list[User]`). Quoted below
+    // so forward / recursive references resolve without ordering constraints.
+    const annotation = node.type as unknown as string;
+    const optional = annotation.startsWith('Optional[');
+    // When the Python attribute differs from the GraphQL wire key (snake_case
+    // conversion or a keyword escaped with `_`), pin the wire key with a pydantic
+    // `Field(alias=...)` so hydration maps it onto the attribute. Optional fields
+    // also default to `None` so a missing key is allowed.
+    let assignment = '';
+    if (attr !== wire) {
+      assignment = optional
+        ? ` = Field(default=None, alias='${wire}')`
+        : ` = Field(alias='${wire}')`;
+    } else if (optional) {
+      assignment = ' = None';
+    }
+    return `${attr}: '${annotation}'${assignment}`;
+  }
+
+  InputValueDefinition(node: InputValueDefinitionNode): string {
+    if (this.config.onlyEnums) return '';
+    return `${this.dictFieldName(node.name.value)}: '${node.type as unknown as string}'`;
+  }
+
+  // --- Class / block builder ---
+
+  private buildClass(name: string, base: string, fields: readonly string[]): string {
+    const header = `class ${name}(${base}):`;
+    const body = fields.filter(Boolean);
+    if (body.length === 0) {
+      return `${header}\n${INDENT}pass`;
+    }
+    return [header, ...body.map(f => INDENT + f)].join('\n');
+  }
+
+  /**
+   * Build a base-model subclass. Fields are class-level annotations; if any field
+   * pins a wire key with `Field(alias=...)`, a
+   * `model_config = ConfigDict(populate_by_name=True)` line is prepended so the
+   * class can be constructed by attribute name as well as by alias.
+   */
+  private buildBaseModelClass(name: string, fields: readonly string[]): string {
+    const header = `class ${name}(${this.config.baseModel}):`;
+    const body = fields.filter(Boolean);
+    if (body.length === 0) {
+      return `${header}\n${INDENT}pass`;
+    }
+    const lines = [header];
+    if (body.some(f => f.includes('alias='))) {
+      lines.push(`${INDENT}model_config = ConfigDict(populate_by_name=True)`);
+    }
+    lines.push(...body.map(f => INDENT + f));
+    return lines.join('\n');
+  }
+
+  /**
+   * Split a rendered `key: 'annotation'` field back into its parts. The key is
+   * always a GraphQL Name (no `: '`) and the annotation is always single-quoted,
+   * so splitting on the first `: '` is unambiguous.
+   */
+  private splitField(rendered: string): { key: string; annotation: string } {
+    const sep = rendered.indexOf(": '");
+    const key = rendered.slice(0, sep);
+    const annotation = rendered.slice(sep + 3, -1);
+    return { key, annotation };
+  }
+
+  /**
+   * Build a `TypedDict`. If every key is a valid Python attribute, the readable
+   * class syntax is used; otherwise (a key is a keyword / non-identifier under
+   * `'keep'`) the functional `TypedDict('Name', { 'key': 'annotation' })` syntax
+   * is emitted so the offending keys remain valid.
+   */
+  private buildTypedDict(name: string, fields: readonly string[]): string {
+    const body = fields.filter(Boolean);
+    const parsed = body.map(f => this.splitField(f));
+    const needsFunctional = parsed.some(p => !isValidClassAttr(p.key));
+    if (!needsFunctional) {
+      return this.buildClass(name, 'TypedDict', body);
+    }
+    const entries = parsed.map(p => `${INDENT}'${p.key}': '${p.annotation}',`);
+    return [`${name} = TypedDict('${name}', {`, ...entries, '})'].join('\n');
+  }
+
+  ObjectTypeDefinition(node: ObjectTypeDefinitionNode): string {
+    if (this.config.onlyOperationTypes || this.config.onlyEnums) return '';
+    if (node.name.value.startsWith('__')) return '';
+    this._usesBaseModel = true;
+    return this.buildBaseModelClass(this.convertName(node), node.fields as unknown as string[]);
+  }
+
+  InterfaceTypeDefinition(node: InterfaceTypeDefinitionNode): string {
+    if (this.config.onlyOperationTypes || this.config.onlyEnums) return '';
+    if (node.name.value.startsWith('__')) return '';
+    this._usesBaseModel = true;
+    return this.buildBaseModelClass(this.convertName(node), node.fields as unknown as string[]);
+  }
+
+  InputObjectTypeDefinition(node: InputObjectTypeDefinitionNode): string {
+    if (this.config.onlyEnums) return '';
+    if (node.name.value.startsWith('__')) return '';
+    this._usesTypedDict = true;
+    // Under `'keep'`, a keyword / non-identifier field key forces functional
+    // `TypedDict` syntax; otherwise readable class syntax is used.
+    return this.buildTypedDict(this.convertName(node), node.fields as unknown as string[]);
+  }
+
+  EnumTypeDefinition(node: EnumTypeDefinitionNode): string {
+    if (node.name.value.startsWith('__')) return '';
+    this._usesEnum = true;
+    const name = this.convertName(node, {
+      useTypesPrefix: this.config.enumPrefix,
+      useTypesSuffix: this.config.enumSuffix,
+    });
+    const members = (node.values ?? []).map(
+      value => `${INDENT}${value.name.value} = '${value.name.value}'`,
+    );
+    const body = members.length > 0 ? members : [`${INDENT}pass`];
+    return [`class ${name}(str, Enum):`, ...body].join('\n');
+  }
+
+  /** Scalars are handled inline via the scalar map; no standalone class. */
+  ScalarTypeDefinition(): string {
+    return '';
+  }
+
+  /**
+   * `union Foo = A | B` -> `Foo = Union['A', 'B']`.
+   *
+   * Must be overridden: `BaseTypesVisitor`'s implementation emits TypeScript
+   * (`export type Foo = A | B;`), and it dereferences the node description
+   * unguarded, so an undocumented union crashes the generator outright. A schema
+   * merely *containing* a union is enough to hit this, even when no operation
+   * selects one.
+   *
+   * Member names come from `parent[key]` — the original, un-visited node — the
+   * same way `typescript`'s visitor reads them. By `leave` time `node.types`
+   * has already been rendered into annotations (`Optional[Foo]`) and would
+   * otherwise be mangled. Members are quoted so a union declared before its
+   * members still resolves, matching how every other annotation in this plugin
+   * is emitted.
+   *
+   * Note this only makes a union *declaration* valid Python. Actually selecting
+   * union members in an operation needs inline-fragment support, which
+   * `python-operations` does not implement yet (it throws).
+   */
+  UnionTypeDefinition(
+    node: UnionTypeDefinitionNode,
+    key: string | number | undefined,
+    parent: any,
+  ): string {
+    if (this.config.onlyOperationTypes || this.config.onlyEnums) return '';
+    if (node.name.value.startsWith('__')) return '';
+
+    const originalNode = parent[key] as UnionTypeDefinitionNode;
+    const members = (originalNode.types ?? []).map(t => `'${this.convertName(t)}'`);
+
+    // A member-less union is legal to declare but has no inhabitant.
+    if (members.length === 0) {
+      return `${this.convertName(node)} = Any`;
+    }
+
+    this._usesUnion = true;
+    return `${this.convertName(node)} = Union[${members.join(', ')}]`;
+  }
+
+  /**
+   * Build the `prepend` import lines, emitting only what the generated
+   * `content` actually references.
+   */
+  getImports(content: string): string[] {
+    const imports: string[] = [];
+
+    if (this._usesEnum) {
+      imports.push('from enum import Enum');
+    }
+
+    const typingParts: string[] = [];
+    if (content.includes('Optional[')) {
+      typingParts.push('Optional');
+    }
+    if (this._usesTypedDict) {
+      typingParts.push('TypedDict');
+    }
+    if (this._usesUnion) {
+      typingParts.push('Union');
+    }
+    if (/\bAny\b/.test(content)) {
+      typingParts.push('Any');
+    }
+    if (typingParts.length > 0) {
+      imports.push(`from typing import ${typingParts.join(', ')}`);
+    }
+
+    if (this._usesBaseModel) {
+      // `ConfigDict` / `Field` are pydantic names; emit them from the same module
+      // as the base model, and only when the generated content uses them.
+      const parts = [this.config.baseModel];
+      if (content.includes('ConfigDict(')) {
+        parts.push('ConfigDict');
+      }
+      if (content.includes('Field(')) {
+        parts.push('Field');
+      }
+      imports.push(`from ${this.config.runtimeImport} import ${parts.join(', ')}`);
+    }
+
+    return imports;
+  }
+}

package/tests/__snapshots__/python.spec.ts.snap

@@ -0,0 +1,37 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`@graphql-codegen/python > matches the full schema snapshot 1`] = `
+"from enum import Enum
+from typing import Optional, TypedDict, Any
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Role(str, Enum):
+    ADMIN = 'ADMIN'
+    USER = 'USER'
+
+
+class Node(BaseModel):
+    id: 'str'
+
+
+class User(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    id: 'str'
+    name: 'str'
+    email: 'Optional[str]' = None
+    role: 'Role'
+    friends: 'list[User]'
+    created_at: 'Optional[Any]' = Field(default=None, alias='createdAt')
+
+
+UserFilter = TypedDict('UserFilter', {
+    'role': 'Optional[Role]',
+    'from': 'Optional[str]',
+})
+
+
+class Query(BaseModel):
+    users: 'list[User]'
+"
+`;

package/tests/python.spec.ts

@@ -0,0 +1,286 @@
+import { buildSchema } from 'graphql';
+import { describe, expect, it } from 'vitest';
+import type { PythonPluginConfig } from '../src/config.js';
+import { plugin } from '../src/index.js';
+
+const FULL_SCHEMA = /* GraphQL */ `
+  scalar DateTime
+  enum Role {
+    ADMIN
+    USER
+  }
+  interface Node {
+    id: ID!
+  }
+  type User implements Node {
+    id: ID!
+    name: String!
+    email: String
+    role: Role!
+    friends: [User!]!
+    createdAt: DateTime
+  }
+  input UserFilter {
+    role: Role
+    from: String
+  }
+  type Query {
+    users(filter: UserFilter): [User!]!
+  }
+`;
+
+function run(schemaSDL: string, config: PythonPluginConfig = {}) {
+  const schema = buildSchema(schemaSDL);
+  const result = plugin(schema, [], config) as { prepend: string[]; content: string };
+  return result;
+}
+
+function full(schemaSDL: string, config: PythonPluginConfig = {}): string {
+  const { prepend, content } = run(schemaSDL, config);
+  return `${prepend.join('\n')}\n\n\n${content}\n`;
+}
+
+describe('@graphql-codegen/python', () => {
+  it('maps built-in scalars to native Python types and strips Optional on non-null', () => {
+    const { content } = run(`
+      type Q { a: ID! b: String! c: Int! d: Float! e: Boolean! }
+    `);
+    expect(content).toContain(`a: 'str'`);
+    expect(content).toContain(`b: 'str'`);
+    expect(content).toContain(`c: 'int'`);
+    expect(content).toContain(`d: 'float'`);
+    expect(content).toContain(`e: 'bool'`);
+  });
+
+  it('renders enums as (str, Enum) classes', () => {
+    const { content, prepend } = run(`enum Role { ADMIN USER } type Q { r: Role! }`);
+    expect(content).toContain('class Role(str, Enum):');
+    expect(content).toContain(`    ADMIN = 'ADMIN'`);
+    expect(content).toContain(`    USER = 'USER'`);
+    expect(prepend).toContain('from enum import Enum');
+  });
+
+  it('renders object types as BaseModel classes', () => {
+    const { content, prepend } = run(`type User { id: ID! name: String! }`);
+    expect(content).toContain('class User(BaseModel):');
+    expect(content).toContain(`    id: 'str'`);
+    expect(content).toContain(`    name: 'str'`);
+    expect(prepend).toContain('from pydantic import BaseModel');
+  });
+
+  it('renders interfaces as BaseModel classes', () => {
+    const { content } = run(`interface Node { id: ID! }`);
+    expect(content).toContain('class Node(BaseModel):');
+    expect(content).toContain(`    id: 'str'`);
+  });
+
+  it('renders input types as TypedDict', () => {
+    const { content, prepend } = run(`input UserFilter { role: String }`);
+    expect(content).toContain('class UserFilter(TypedDict):');
+    expect(prepend).toContain('from typing import Optional, TypedDict');
+  });
+
+  it('maps custom scalars to Any', () => {
+    const { content, prepend } = run(`scalar DateTime type Q { at: DateTime }`);
+    expect(content).toContain(`at: 'Optional[Any]'`);
+    // no standalone class for the scalar
+    expect(content).not.toContain('class DateTime');
+    expect(prepend.some(l => l.includes('Any'))).toBe(true);
+  });
+
+  it('honors config.scalars overrides for custom scalars', () => {
+    const { content } = run(`scalar DateTime type Q { at: DateTime! }`, {
+      scalars: { DateTime: 'datetime' },
+    });
+    expect(content).toContain(`at: 'datetime'`);
+  });
+
+  it('converts object field names to snake_case by default (namingConvention: snakeCase)', () => {
+    const { content } = run(`type User { createdAt: String firstName: String! }`);
+    expect(content).toContain(`created_at: 'Optional[str]'`);
+    expect(content).toContain(`first_name: 'str'`);
+  });
+
+  it('emits a functional TypedDict when an input key is a reserved keyword (inputs keep wire names)', () => {
+    const { content, prepend } = run(`input F { from: String role: String }`);
+    // `from` is a Python keyword -> whole input type uses functional syntax.
+    expect(content).toContain(`F = TypedDict('F', {`);
+    expect(content).toContain(`    'from': 'Optional[str]',`);
+    expect(content).toContain(`    'role': 'Optional[str]',`);
+    expect(content).not.toContain(`class F(TypedDict):`);
+    expect(prepend).toContain('from typing import Optional, TypedDict');
+  });
+
+  it('keeps input keys as wire names (never snake-cased), class syntax when valid', () => {
+    const { content } = run(`input UserFilter { role: String createdAt: String }`);
+    expect(content).toContain(`class UserFilter(TypedDict):`);
+    expect(content).toContain(`    role: 'Optional[str]'`);
+    // Input keys keep the WIRE name (the runtime sends the dict as-is).
+    expect(content).toContain(`    createdAt: 'Optional[str]'`);
+    expect(content).not.toContain(`TypedDict('UserFilter'`);
+  });
+
+  it('input keys keep wire names even under namingConvention: snakeCase', () => {
+    const { content } = run(`input F { from: String createdAt: String }`, {
+      namingConvention: 'snakeCase',
+    });
+    expect(content).toContain(`F = TypedDict('F', {`);
+    expect(content).toContain(`    'from': 'Optional[str]',`);
+    expect(content).toContain(`    'createdAt': 'Optional[str]',`);
+  });
+
+  it('escapes keyword fields on object types and pins the wire key with Field(alias=...)', () => {
+    // Object types extend a pydantic BaseModel; a keyword field is escaped with
+    // `_` and `Field(alias=...)` pins the original wire key so pydantic hydrates
+    // it onto the escaped attribute. `model_config` enables population by name.
+    const { content, prepend } = run(
+      `type Message { from: String! createdAt: String! class: Int }`,
+    );
+    expect(content).toContain(`class Message(BaseModel):`);
+    expect(content).toContain(`    model_config = ConfigDict(populate_by_name=True)`);
+    expect(content).toContain(`    from_: 'str' = Field(alias='from')`);
+    expect(content).toContain(`    created_at: 'str' = Field(alias='createdAt')`);
+    expect(content).toContain(`    class_: 'Optional[int]' = Field(default=None, alias='class')`);
+    expect(content).not.toContain(`def __init__`);
+    expect(prepend).toContain('from pydantic import BaseModel, ConfigDict, Field');
+  });
+
+  it('namingConvention: keep leaves object field names verbatim (keyword -> _ + alias)', () => {
+    const { content } = run(`type User { createdAt: String firstName: String! from: String }`, {
+      namingConvention: 'keep',
+    });
+    expect(content).toContain(`createdAt: 'Optional[str]' = None`);
+    expect(content).toContain(`firstName: 'str'`);
+    expect(content).toContain(`from_: 'Optional[str]' = Field(default=None, alias='from')`);
+    expect(content).not.toContain(`def __init__`);
+  });
+
+  it('composes nullability and lists correctly', () => {
+    const { content } = run(`
+      type User { id: ID! }
+      type Q {
+        nn: String!
+        opt: String
+        reqList: [User!]!
+        nullListNullItem: [User]
+      }
+    `);
+    expect(content).toContain(`nn: 'str'`);
+    expect(content).toContain(`opt: 'Optional[str]'`);
+    expect(content).toContain(`req_list: 'list[User]'`);
+    expect(content).toContain(`null_list_null_item: 'Optional[list[Optional[User]]]'`);
+  });
+
+  it('verifies onlyOperationTypes emits enums + inputs + scalars and skips object types', () => {
+    const { content } = run(
+      `scalar DateTime
+       enum Role { ADMIN USER }
+       input UserFilter { from: String role: Role }
+       type Message { from: String! createdAt: DateTime! }
+       type Query { m: Message }`,
+      { onlyOperationTypes: true },
+    );
+    // enums + input types present...
+    expect(content).toContain(`class Role(str, Enum):`);
+    expect(content).toContain(`UserFilter = TypedDict('UserFilter', {`);
+    // ...object types skipped.
+    expect(content).not.toContain(`class Message`);
+    expect(content).not.toContain(`class Query`);
+    expect(content).not.toContain('BaseModel');
+  });
+
+  it('supports recursive types via quoted forward references', () => {
+    const { content } = run(`type User { id: ID! friends: [User!]! }`);
+    expect(content).toContain('class User(BaseModel):');
+    expect(content).toContain(`friends: 'list[User]'`);
+  });
+
+  it('honors baseModel and runtimeImport config', () => {
+    const { content, prepend } = run(`type User { id: ID! }`, {
+      baseModel: 'PydModel',
+      runtimeImport: 'myapp.models',
+    });
+    expect(content).toContain('class User(PydModel):');
+    expect(prepend).toContain('from myapp.models import PydModel');
+  });
+
+  it('only emits imports that are actually used', () => {
+    // No enum, no input, no Any -> only Optional + BaseModel
+    const { prepend } = run(`type User { name: String }`);
+    expect(prepend).not.toContain('from enum import Enum');
+    expect(prepend).toContain('from typing import Optional');
+    expect(prepend).toContain('from pydantic import BaseModel');
+  });
+
+  it('does not emit introspection types', () => {
+    const { content } = run(FULL_SCHEMA);
+    expect(content).not.toContain('__Schema');
+    expect(content).not.toContain('__Type');
+  });
+
+  describe('unions', () => {
+    const UNION_SCHEMA = /* GraphQL */ `
+      type A {
+        a: String
+      }
+      type B {
+        b: Int
+      }
+      union AB = A | B
+      type Query {
+        thing: AB
+      }
+    `;
+
+    it('emits a typing.Union alias with quoted members', () => {
+      const { content, prepend } = run(UNION_SCHEMA);
+      expect(content).toContain(`Ab = Union['A', 'B']`);
+      expect(prepend.some(l => /from typing import .*\bUnion\b/.test(l))).toBe(true);
+    });
+
+    it('never emits TypeScript for a union', () => {
+      // The inherited BaseTypesVisitor implementation emits `export type Ab = A | B;`
+      // plus a `/** ... */` comment — neither is valid Python.
+      const { content } = run(UNION_SCHEMA);
+      expect(content).not.toContain('export type');
+      expect(content).not.toContain('/**');
+    });
+
+    it('does not crash on a union without a description', () => {
+      // The inherited implementation dereferences the description unguarded.
+      expect(() => run(UNION_SCHEMA)).not.toThrow();
+    });
+
+    it('handles a documented union the same way', () => {
+      const { content } = run(/* GraphQL */ `
+        type A {
+          a: String
+        }
+        type B {
+          b: Int
+        }
+        "Documented."
+        union AB = A | B
+        type Query {
+          thing: AB
+        }
+      `);
+      expect(content).toContain(`Ab = Union['A', 'B']`);
+      expect(content).not.toContain('export type');
+    });
+
+    it('is skipped under onlyOperationTypes / onlyEnums', () => {
+      expect(run(UNION_SCHEMA, { onlyOperationTypes: true }).content).not.toContain('Ab = Union');
+      expect(run(UNION_SCHEMA, { onlyEnums: true }).content).not.toContain('Ab = Union');
+    });
+
+    it('references the union from a field that selects it', () => {
+      const { content } = run(UNION_SCHEMA);
+      expect(content).toContain(`thing: 'Optional[Ab]'`);
+    });
+  });
+
+  it('matches the full schema snapshot', () => {
+    expect(full(FULL_SCHEMA)).toMatchSnapshot();
+  });
+});

package/vitest.config.ts

@@ -0,0 +1,12 @@
+import { defineProject, mergeConfig } from 'vitest/config';
+import { sharedConfig } from '../../../../vitest.config.js';
+
+export default mergeConfig(
+  sharedConfig,
+  defineProject({
+    test: {
+      name: 'python',
+      include: ['**/*.spec.ts'],
+    },
+  }),
+);