schema

type DefaultKind ¶

type DefaultKind int

DefaultKind tags the active field in DefaultValue.

const (
	DefaultNone DefaultKind = iota
	DefaultBool
	DefaultInt
	DefaultFloat
	DefaultString
	DefaultList
	DefaultDict
	DefaultSet
)

type DefaultValue ¶

type DefaultValue struct {
	Kind     DefaultKind
	Bool     bool
	Int      int64
	Float    float64
	Str      string
	List     []DefaultValue
	DictKeys []DefaultValue // parallel with DictVals
	DictVals []DefaultValue
}

DefaultValue represents a statically-parsed Python literal.

func (d DefaultValue) MarshalJSON() ([]byte, error)

func (d DefaultValue) ToJSON() any

type FieldType ¶

type FieldType struct {
	Primitive  PrimitiveType
	Repetition Repetition
}

FieldType combines a primitive type with its cardinality.

func ResolveFieldType(ann TypeAnnotation, ctx *ImportContext) (FieldType, error)

func (ft FieldType) JSONType() map[string]any

type ImportContext ¶

type ImportContext struct {
	// Names maps local name → (module, original_name)
	Names *OrderedMap[string, ImportEntry]
}

ImportContext tracks what names are imported from which modules.

func NewImportContext() *ImportContext

func (ctx *ImportContext) IsBaseModel(name string) bool

func (ctx *ImportContext) IsBasePredictor(name string) bool

func (ctx *ImportContext) IsCogType(name string) bool

func (ctx *ImportContext) IsTypingType(name string) bool

type ImportEntry ¶

type ImportEntry struct {
	Module   string
	Original string
}

ImportEntry records where a name was imported from.

type InputField ¶

type InputField struct {
	Name        string
	Order       int
	FieldType   FieldType
	Default     *DefaultValue
	Description *string
	GE          *float64
	LE          *float64
	MinLength   *uint64
	MaxLength   *uint64
	Regex       *string
	Choices     []DefaultValue
	Deprecated  *bool
}

InputField represents one parameter of predict/train.

func (f *InputField) IsRequired() bool

type Mode ¶

type Mode int

Mode selects whether to extract predict or train signatures.

const (
	ModePredict Mode = iota
	ModeTrain
)

type ModelClassMap ¶

type ModelClassMap = *OrderedMap[string, []ModelField]

ModelClassMap maps class names to their fields.

type ModelField ¶

type ModelField struct {
	Name    string
	Type    TypeAnnotation
	Default *DefaultValue
}

ModelField is a field extracted from a BaseModel subclass.

type OrderedMap ¶

type OrderedMap[K comparable, V any] struct {
	// contains filtered or unexported fields
}

OrderedMap is a simple insertion-ordered map.

func NewOrderedMap[K comparable, V any]() *OrderedMap[K, V]

func (m *OrderedMap[K, V]) Entries(fn func(key K, value V))

func (m *OrderedMap[K, V]) Get(key K) (V, bool)

func (m *OrderedMap[K, V]) Keys() []K

func (m *OrderedMap[K, V]) Len() int

func (m *OrderedMap[K, V]) Set(key K, value V)

type Parser ¶

type Parser func(source []byte, predictRef string, mode Mode, sourceDir string) (*PredictorInfo, error)

Parser is a function that parses source code and extracts predictor info. This is defined as a type to avoid an import cycle between schema and schema/python. The concrete implementation is python.ParsePredictor.

sourceDir is the project root directory, used for resolving cross-file imports (e.g. "from .types import Output"). Pass "" if unknown.

type PredictorInfo ¶

type PredictorInfo struct {
	Inputs *OrderedMap[string, InputField]
	Output SchemaType
	Mode   Mode
}

PredictorInfo is the top-level extraction result.

type PrimitiveType ¶

type PrimitiveType int

PrimitiveType maps Python types to JSON Schema types.

const (
	TypeBool PrimitiveType = iota
	TypeFloat
	TypeInteger
	TypeString
	TypePath   // cog.Path — {"type":"string","format":"uri"}
	TypeFile   // cog.File (deprecated) — same wire format as Path
	TypeSecret // cog.Secret — write-only, masked
	TypeAny    // typing.Any or unresolved
)

func PrimitiveFromName(name string) (PrimitiveType, bool)

func (p PrimitiveType) JSONType() map[string]any

func (p PrimitiveType) String() string

type Repetition ¶

type Repetition int

Repetition describes cardinality of a field.

const (
	Required Repetition = iota
	Optional
	Repeated // list[X]
)

type SchemaError ¶

type SchemaError struct {
	Kind    SchemaErrorKind
	Message string
}

SchemaError represents errors during schema generation.

func NewError(kind SchemaErrorKind, msg string) *SchemaError

func WrapError(kind SchemaErrorKind, msg string, inner error) *SchemaError

func (e *SchemaError) Error() string

type SchemaErrorKind ¶

type SchemaErrorKind int

SchemaErrorKind classifies schema generation errors.

const (
	ErrParse SchemaErrorKind = iota
	ErrPredictorNotFound
	ErrMethodNotFound
	ErrMissingReturnType
	ErrMissingTypeAnnotation
	ErrUnsupportedType
	ErrDefaultFactoryNotSupported
	ErrInvalidConstraint
	ErrInvalidPredictRef
	ErrOptionalOutput
	ErrConcatIteratorNotStr
	ErrChoicesNotResolvable
	ErrDefaultNotResolvable
	ErrUnresolvableType
	ErrOther
)

type SchemaField ¶

type SchemaField struct {
	Type     SchemaType
	Default  *DefaultValue
	Required bool
}

SchemaField is a named field within a SchemaObject.

type SchemaType ¶

type SchemaType struct {
	Kind SchemaTypeKind

	// Primitive: for Kind=SchemaPrimitive — one of the base scalar types.
	Primitive PrimitiveType

	// Array: for Kind=SchemaArray — the element type.
	Items *SchemaType

	// Dict: for Kind=SchemaDict — key and value types.
	// KeyType is always string in JSON Schema, but we track it for completeness.
	KeyType   *SchemaType
	ValueType *SchemaType

	// Object: for Kind=SchemaObject — named fields with types and defaults.
	Fields *OrderedMap[string, SchemaField]

	// Iterator/ConcatIterator: for Kind=SchemaIterator|SchemaConcatIterator.
	// The yielded element type.
	Elem *SchemaType

	// Nullable: wraps any type to allow null.
	Nullable bool
}

SchemaType is a recursive algebraic data type representing any type that can appear in a Cog predictor's output (or, in the future, input) position.

It replaces the flat OutputType/PrimitiveType system with a composable type tree that can represent dict[str, list[int]], nested BaseModel subclasses, TypedDicts, and types resolved from .pyi stubs — all without running Python.

func ResolveSchemaType(ann TypeAnnotation, ctx *ImportContext, models ModelClassMap) (SchemaType, error)

dict[str, list[dict[str, int]]]  →  SchemaDict{ValueType: SchemaArray{Items: SchemaDict{...}}}
list[dict[str, str]]             →  SchemaArray{Items: SchemaDict{ValueType: SchemaPrim(TypeString)}}

func SchemaAnyType() SchemaType

func SchemaArrayOf(elem SchemaType) SchemaType

func SchemaConcatIteratorOf() SchemaType

func SchemaDictOf(value SchemaType) SchemaType

func SchemaIteratorOf(elem SchemaType) SchemaType

func SchemaObjectOf(fields *OrderedMap[string, SchemaField]) SchemaType

func SchemaPrim(p PrimitiveType) SchemaType

func (s SchemaType) JSONSchema() map[string]any

type SchemaTypeKind ¶

type SchemaTypeKind int

SchemaTypeKind tags the active variant in SchemaType.

const (
	// SchemaPrimitive is a scalar type: bool, int, float, str, Path, File, Secret.
	SchemaPrimitive SchemaTypeKind = iota
	// SchemaAny is an opaque JSON value (unparameterized dict, Any, etc).
	SchemaAny
	// SchemaArray is a homogeneous list/array.
	SchemaArray
	// SchemaDict is a string-keyed dictionary with a typed value.
	SchemaDict
	// SchemaObject is a product type with named fields (BaseModel, TypedDict, dataclass).
	SchemaObject
	// SchemaIterator is a cog Iterator[T] — array with x-cog-array-type=iterator.
	SchemaIterator
	// SchemaConcatIterator is a cog ConcatenateIterator[str] — streaming text.
	SchemaConcatIterator
)

type TypeAnnotation ¶

type TypeAnnotation struct {
	Kind TypeAnnotationKind
	Name string           // for Simple
	Args []TypeAnnotation // for Generic (outer=Name, args=Args) or Union (members=Args)
}

TypeAnnotation is a parsed Python type annotation (intermediate, before resolution).

func UnwrapOptional(ann TypeAnnotation) (TypeAnnotation, bool)

type TypeAnnotationKind ¶

type TypeAnnotationKind int

TypeAnnotationKind tags the variant.

const (
	TypeAnnotSimple TypeAnnotationKind = iota
	TypeAnnotGeneric
	TypeAnnotUnion
)