openapi-ts / packages / openapi-python / src / plugins / pydantic / v2 / plugin.ts

import type { SymbolMeta } from '@hey-api/codegen-core';
import { fromRef, ref } from '@hey-api/codegen-core';
import type { IR, SchemaWithType } from '@hey-api/shared';
import { deduplicateSchema, pathToJsonPointer } from '@hey-api/shared';

// import { $ } from '../../../py-dsl';
import type { Ast, IrSchemaToAstOptions } from '../shared/types';
import type { PydanticPlugin } from '../types';
import { createProcessor } from './processor';
import { irSchemaWithTypeToAst } from './toAst';

export function irSchemaToAst({
  optional,
  plugin,
  schema,
  schemaExtractor,
  state,
}: IrSchemaToAstOptions & {
  optional?: boolean;
  schema: IR.SchemaObject;
}): Ast {
  if (schemaExtractor && !schema.$ref) {
    const extracted = schemaExtractor({
      meta: {
        resource: 'definition',
        resourceId: pathToJsonPointer(fromRef(state.path)),
      },
      naming: plugin.config.definitions,
      path: fromRef(state.path),
      plugin,
      schema,
    });
    if (extracted !== schema) schema = extracted;
  }

  if (schema.$ref) {
    const query: SymbolMeta = {
      category: 'schema',
      resource: 'definition',
      resourceId: schema.$ref,
      tool: 'pydantic',
    };
    const refSymbol = plugin.referenceSymbol(query);
    const refName = typeof refSymbol === 'string' ? refSymbol : refSymbol.name;

    return {
      // expression: $.expr(refName),
      fieldConstraints: optional ? { default: null } : undefined,
      hasLazyExpression: !plugin.isSymbolRegistered(query),
      models: [],
      // pipes: [],
      typeAnnotation: refName,
    };
  }

  if (schema.type) {
    const typeAst = irSchemaWithTypeToAst({
      plugin,
      schema: schema as SchemaWithType,
      state,
    });

    const constraints: Record<string, unknown> = {};
    if (optional) {
      constraints.default = null;
    }
    if (schema.default !== undefined) {
      constraints.default = schema.default;
    }
    if (schema.description) {
      constraints.description = schema.description;
    }

    return {
      ...typeAst,
      fieldConstraints: { ...typeAst.fieldConstraints, ...constraints },
      // pipes: [],
    };
  }

  if (schema.items) {
    schema = deduplicateSchema({ schema });

    if (schema.items) {
      const itemsAnnotations: string[] = [];
      const itemsConstraints: Record<string, unknown>[] = [];

      for (const item of schema.items) {
        const itemAst = irSchemaToAst({
          plugin,
          schema: item,
          state: {
            ...state,
            path: ref([...fromRef(state.path), 'items']),
          },
        });
        itemsAnnotations.push(itemAst.typeAnnotation);
        if (itemAst.fieldConstraints) {
          itemsConstraints.push(itemAst.fieldConstraints);
        }
      }

      const unionType = itemsAnnotations.join(' | ');
      return {
        // expression: $.expr(`list[${unionType}]`),
        fieldConstraints: itemsConstraints.length > 0 ? itemsConstraints[0] : undefined,
        hasLazyExpression: false,
        models: [],
        // pipes: [],
        typeAnnotation: `list[${unionType}]`,
      };
    }
  }

  return {
    // expression: $.expr('Any'),
    hasLazyExpression: false,
    models: [],
    // pipes: [],
    typeAnnotation: 'Any',
  };
}

export const handlerV2: PydanticPlugin['Handler'] = ({ plugin }) => {
  plugin.symbol('Any', {
    external: 'typing',
    importKind: 'named',
    meta: {
      category: 'external',
      resource: 'typing.Any',
    },
  });
  plugin.symbol('BaseModel', {
    external: 'pydantic',
    importKind: 'named',
    meta: {
      category: 'external',
      resource: 'pydantic.BaseModel',
    },
  });
  plugin.symbol('ConfigDict', {
    external: 'pydantic',
    importKind: 'named',
    meta: {
      category: 'external',
      resource: 'pydantic.ConfigDict',
    },
  });
  plugin.symbol('Field', {
    external: 'pydantic',
    importKind: 'named',
    meta: {
      category: 'external',
      resource: 'pydantic.Field',
    },
  });
  plugin.symbol('Literal', {
    external: 'typing',
    importKind: 'named',
    meta: {
      category: 'external',
      resource: 'typing.Literal',
    },
  });
  plugin.symbol('Optional', {
    external: 'typing',
    importKind: 'named',
    meta: {
      category: 'external',
      resource: 'typing.Optional',
    },
  });

  const processor = createProcessor(plugin);

  plugin.forEach('operation', 'parameter', 'requestBody', 'schema', 'webhook', (event) => {
    switch (event.type) {
      case 'parameter':
        processor.process({
          meta: {
            resource: 'definition',
            resourceId: pathToJsonPointer(event._path),
          },
          naming: plugin.config.definitions,
          path: event._path,
          plugin,
          schema: event.parameter.schema,
          tags: event.tags,
        });
        break;
      case 'requestBody':
        processor.process({
          meta: {
            resource: 'definition',
            resourceId: pathToJsonPointer(event._path),
          },
          naming: plugin.config.definitions,
          path: event._path,
          plugin,
          schema: event.requestBody.schema,
          tags: event.tags,
        });
        break;
      case 'schema':
        processor.process({
          meta: {
            resource: 'definition',
            resourceId: pathToJsonPointer(event._path),
          },
          naming: plugin.config.definitions,
          path: event._path,
          plugin,
          schema: event.schema,
          tags: event.tags,
        });
        break;
    }
  });
};