feat: add resolvers to pydantic

Author: mrlubos

dev/python/presets.ts

@@ -9,11 +9,14 @@ export const presets = {
         containerName: 'OpenCode',
         strategy: 'single',
       },
+      paramsStructure: 'flat',
     }),
   ],
   validated: () => [
     /** SDK + Pydantic validation */
-    sdk(),
+    sdk({
+      paramsStructure: 'flat',
+    }),
     pydantic(),
   ],
 } as const satisfies Record<string, () => ReadonlyArray<PluginConfig>>;

packages/openapi-python/src/plugins/pydantic/index.ts

@@ -1,2 +1,9 @@
 export { defaultConfig, defineConfig } from './config';
+export type {
+  EnumResolverContext,
+  NumberResolverContext,
+  ObjectResolverContext,
+  Resolvers,
+  StringResolverContext,
+} from './resolvers';
 export type { PydanticPlugin } from './types';

packages/openapi-python/src/plugins/pydantic/resolvers.ts

@@ -0,0 +1,99 @@
+import type { Plugin, SchemaVisitorContext, SchemaWithType, Walker } from '@hey-api/shared';
+
+import type { DollarPyDsl } from '../../py-dsl';
+import type { PydanticField, PydanticFinal, PydanticResult, PydanticType } from './shared/types';
+import type { PydanticPlugin } from './types';
+
+export type Resolvers = Plugin.Resolvers<{
+  /**
+   * Resolver for enum schemas.
+   *
+   * Allows customization of how enum types are rendered.
+   *
+   * Returning `undefined` will execute the default resolver logic.
+   */
+  enum?: (ctx: EnumResolverContext) => PydanticType | undefined;
+  /**
+   * Resolver for number schemas.
+   *
+   * Allows customization of how number types are rendered.
+   *
+   * Returning `undefined` will execute the default resolver logic.
+   */
+  number?: (ctx: NumberResolverContext) => PydanticType | undefined;
+  /**
+   * Resolver for object schemas.
+   *
+   * Allows customization of how object types are rendered.
+   *
+   * Returning `undefined` will execute the default resolver logic.
+   */
+  object?: (
+    ctx: ObjectResolverContext,
+  ) => (PydanticType & { fields?: Array<PydanticField> }) | undefined;
+  /**
+   * Resolver for string schemas.
+   *
+   * Allows customization of how string types are rendered.
+   *
+   * Returning `undefined` will execute the default resolver logic.
+   */
+  string?: (ctx: StringResolverContext) => PydanticType | undefined;
+}>;
+
+interface BaseContext extends DollarPyDsl {
+  /** The plugin instance. */
+  plugin: PydanticPlugin['Instance'];
+}
+
+export interface EnumResolverContext extends BaseContext {
+  /**
+   * Nodes used to build different parts of the result.
+   */
+  nodes: {
+    base: (ctx: EnumResolverContext) => PydanticType;
+    items: (ctx: EnumResolverContext) => {
+      enumMembers: Required<PydanticFinal>['enumMembers'];
+      isNullable: boolean;
+    };
+  };
+  schema: SchemaWithType<'enum'>;
+}
+
+export interface NumberResolverContext extends BaseContext {
+  /**
+   * Nodes used to build different parts of the result.
+   */
+  nodes: {
+    base: (ctx: NumberResolverContext) => PydanticType;
+    const: (ctx: NumberResolverContext) => PydanticType | undefined;
+  };
+  schema: SchemaWithType<'integer' | 'number'>;
+}
+
+export interface ObjectResolverContext extends BaseContext {
+  _childResults: Array<PydanticResult>;
+  applyModifiers: (result: PydanticResult, opts: { optional?: boolean }) => PydanticFinal;
+  /**
+   * Nodes used to build different parts of the result.
+   */
+  nodes: {
+    additionalProperties: (ctx: ObjectResolverContext) => PydanticType | null | undefined;
+    base: (ctx: ObjectResolverContext) => PydanticType & { fields?: Array<PydanticField> };
+    fields: (ctx: ObjectResolverContext) => Array<PydanticField>;
+  };
+  schema: SchemaWithType<'object'>;
+  walk: Walker<PydanticResult, PydanticPlugin['Instance']>;
+  walkerCtx: SchemaVisitorContext<PydanticPlugin['Instance']>;
+}
+
+export interface StringResolverContext extends BaseContext {
+  /**
+   * Nodes used to build different parts of the result.
+   */
+  nodes: {
+    base: (ctx: StringResolverContext) => PydanticType;
+    const: (ctx: StringResolverContext) => PydanticType | undefined;
+  };
+  schema: SchemaWithType<'string'>;
+}

packages/openapi-python/src/plugins/pydantic/types.ts

@@ -7,10 +7,13 @@ import type {
   Plugin,
 } from '@hey-api/shared';
 
+import type { Resolvers } from './resolvers';
+
 export type UserConfig = Plugin.Name<'pydantic'> &
   Plugin.Hooks &
   Plugin.UserComments &
-  Plugin.UserExports & {
+  Plugin.UserExports &
+  Resolvers & {
     /**
      * Casing convention for generated names.
      *
@@ -185,7 +188,8 @@ export type UserConfig = Plugin.Name<'pydantic'> &
 export type Config = Plugin.Name<'pydantic'> &
   Plugin.Hooks &
   Plugin.Comments &
-  Plugin.Exports & {
+  Plugin.Exports &
+  Resolvers & {
     /** Casing convention for generated names. */
     case: Casing;
     /** Configuration for reusable schema definitions. */

packages/openapi-python/src/plugins/pydantic/v2/toAst/enum.ts

@@ -1,8 +1,8 @@
-import type { Symbol } from '@hey-api/codegen-core';
 import type { SchemaWithType } from '@hey-api/shared';
 import { toCase } from '@hey-api/shared';
 
 import { $ } from '../../../../py-dsl';
+import type { EnumResolverContext } from '../../resolvers';
 import type { PydanticFinal, PydanticType } from '../../shared/types';
 import type { PydanticPlugin } from '../../types';
 
@@ -21,13 +21,8 @@ function toEnumMemberName(value: string | number): string {
   return toCase(value, 'SCREAMING_SNAKE_CASE');
 }
 
-function extractEnumMembers(
-  schema: SchemaWithType<'enum'>,
-  plugin: PydanticPlugin['Instance'],
-): {
-  enumMembers: Required<PydanticFinal>['enumMembers'];
-  isNullable: boolean;
-} {
+function itemsNode(ctx: EnumResolverContext) {
+  const { plugin, schema } = ctx;
   const enumMembers: Required<PydanticFinal>['enumMembers'] = [];
   let isNullable = false;
 
@@ -51,51 +46,67 @@ function extractEnumMembers(
   return { enumMembers, isNullable };
 }
 
-function toLiteralType(
-  enumMembers: Required<PydanticFinal>['enumMembers'],
-  plugin: PydanticPlugin['Instance'],
-): string | Symbol | ReturnType<typeof $.subscript> {
+function baseNode(ctx: EnumResolverContext): PydanticType {
+  const { plugin } = ctx;
+  const { enumMembers } = ctx.nodes.items(ctx);
+
   if (enumMembers.length === 0) {
-    return plugin.external('typing.Any');
+    return {
+      type: plugin.external('typing.Any'),
+    };
+  }
+
+  const mode = plugin.config.enums ?? 'enum';
+
+  if (mode === 'literal') {
+    if (enumMembers.length === 0) {
+      return {
+        type: plugin.external('typing.Any'),
+      };
+    }
+
+    const literal = plugin.external('typing.Literal');
+    const values = enumMembers.map((m) =>
+      // TODO: replace
+      typeof m.value === 'string' ? `"<<<<${m.value}"` : `<<<${m.value}`,
+    );
+
+    return {
+      type: $(literal).slice(...values),
+    };
   }
 
-  const literal = plugin.external('typing.Literal');
-  const values = enumMembers.map((m) =>
-    // TODO: replace
-    typeof m.value === 'string' ? `"<<<<${m.value}"` : `<<<${m.value}`,
-  );
+  return {};
+}
 
-  return $(literal).slice(...values);
+function enumResolver(ctx: EnumResolverContext): PydanticType {
+  return ctx.nodes.base(ctx);
 }
 
 export function enumToType({
-  mode = 'enum',
   plugin,
   schema,
 }: {
-  mode?: 'enum' | 'literal';
   plugin: PydanticPlugin['Instance'];
   schema: SchemaWithType<'enum'>;
 }): EnumToTypeResult {
-  const { enumMembers, isNullable } = extractEnumMembers(schema, plugin);
+  const ctx: EnumResolverContext = {
+    $,
+    nodes: {
+      base: baseNode,
+      items: itemsNode,
+    },
+    plugin,
+    schema,
+  };
 
-  if (enumMembers.length === 0) {
-    return {
-      enumMembers,
-      isNullable,
-      type: plugin.external('typing.Any'),
-    };
-  }
+  const resolver = plugin.config['~resolvers']?.enum;
+  const resolved = resolver?.(ctx) ?? enumResolver(ctx);
 
-  if (mode === 'literal') {
-    return {
-      enumMembers,
-      isNullable,
-      type: toLiteralType(enumMembers, plugin),
-    };
-  }
+  const { enumMembers, isNullable } = ctx.nodes.items(ctx);
 
   return {
+    ...resolved,
     enumMembers,
     isNullable,
   };

packages/openapi-python/src/plugins/pydantic/v2/toAst/number.ts

@@ -1,18 +1,13 @@
 import type { SchemaWithType } from '@hey-api/shared';
 
 import { $ } from '../../../../py-dsl';
+import type { NumberResolverContext } from '../../resolvers';
 import type { PydanticType } from '../../shared/types';
 import type { PydanticPlugin } from '../../types';
 import type { FieldConstraints } from '../constants';
 
-export function numberToType({
-  plugin,
-  schema,
-}: {
-  plugin: PydanticPlugin['Instance'];
-  schema: SchemaWithType<'integer' | 'number'>;
-}): PydanticType {
-  const constraints: FieldConstraints = {};
+function constNode(ctx: NumberResolverContext): PydanticType | undefined {
+  const { plugin, schema } = ctx;
 
   if (typeof schema.const === 'number') {
     const literal = plugin.external('typing.Literal');
@@ -21,6 +16,14 @@ export function numberToType({
     };
   }
 
+  return undefined;
+}
+
+function baseNode(ctx: NumberResolverContext): PydanticType {
+  const { schema } = ctx;
+
+  const constraints: FieldConstraints = {};
+
   if (schema.minimum !== undefined) {
     constraints.ge = schema.minimum;
   }
@@ -46,3 +49,31 @@ export function numberToType({
     type: schema.type === 'integer' ? 'int' : 'float',
   };
 }
+
+function numberResolver(ctx: NumberResolverContext): PydanticType {
+  const constResult = ctx.nodes.const(ctx);
+  if (constResult) return constResult;
+
+  return ctx.nodes.base(ctx);
+}
+
+export function numberToType({
+  plugin,
+  schema,
+}: {
+  plugin: PydanticPlugin['Instance'];
+  schema: SchemaWithType<'integer' | 'number'>;
+}): PydanticType {
+  const ctx: NumberResolverContext = {
+    $,
+    nodes: {
+      base: baseNode,
+      const: constNode,
+    },
+    plugin,
+    schema,
+  };
+
+  const resolver = plugin.config['~resolvers']?.number;
+  return resolver?.(ctx) ?? numberResolver(ctx);
+}