zzet--gortex/internal/contracts/shape_py.go

package contracts

import (
	"regexp"
	"strings"
)

// pyFieldRe captures a Python class attribute with a type annotation.
// Pydantic V1/V2 models and plain typed dataclasses both use this
// form. Capture groups:
//
//	1: attribute name
//	2: type expression (up to the optional `=` or end of line)
//	3: default value expression (after `=`) — empty when no default
var pyFieldRe = regexp.MustCompile(`^\s*([A-Za-z_][\w]*)\s*:\s*([^=#\n]+?)(?:\s*=\s*(.+?))?\s*(?:#\s*(.*))?$`)

// pyFieldAliasRe pulls alias="wire" out of a Pydantic Field(...) call,
// which is how users rename a Python-legal attribute to a JSON-valid
// wire name.
var pyFieldAliasRe = regexp.MustCompile(`alias\s*=\s*['"]([^'"]+)['"]`)

// extractPythonShape reads the body of a class definition and returns
// its typed attributes. Attributes without a type annotation are
// skipped because we can't record a useful type for them. Methods
// are skipped by requiring `:` before any `(`.
func extractPythonShape(src []byte, startLine, endLine int) *Shape {
	body := sliceBody(src, startLine, endLine)
	if body == "" {
		return nil
	}
	lines := strings.Split(body, "\n")
	if len(lines) == 0 {
		return nil
	}

	// Determine the indent of the class body so we can stop scanning
	// at the next declaration at or below the class indent. Python
	// classes have `class Foo(Bar):` on the first line and the body
	// on subsequent indented lines.
	classIndent := leadingWhitespace(lines[0])
	shape := &Shape{Kind: "class"}

	for i, raw := range lines {
		if i == 0 {
			continue // class header
		}
		if strings.TrimSpace(raw) == "" {
			continue
		}
		indent := leadingWhitespace(raw)
		if len(indent) <= len(classIndent) {
			// De-dent — end of class body.
			break
		}
		line := strings.TrimSpace(raw)
		if strings.HasPrefix(line, "#") {
			continue
		}
		// Skip methods: def / async def.
		if strings.HasPrefix(line, "def ") || strings.HasPrefix(line, "async def ") {
			continue
		}
		// Skip decorator lines at field scope.
		if strings.HasPrefix(line, "@") {
			continue
		}
		// Require a `:` to appear before any `(`.
		colonIdx := strings.Index(line, ":")
		if colonIdx < 0 {
			continue
		}
		parenIdx := strings.Index(line, "(")
		if parenIdx >= 0 && parenIdx < colonIdx {
			continue
		}
		m := pyFieldRe.FindStringSubmatch(line)
		if m == nil {
			continue
		}
		name := m[1]
		typeExpr := strings.TrimSpace(m[2])
		defaultExpr := ""
		comment := ""
		if len(m) > 3 {
			defaultExpr = strings.TrimSpace(m[3])
		}
		if len(m) > 4 {
			comment = truncateComment(m[4])
		}
		if typeExpr == "" {
			continue
		}
		alias := ""
		if defaultExpr != "" {
			if am := pyFieldAliasRe.FindStringSubmatch(defaultExpr); len(am) > 1 {
				alias = am[1]
			}
		}
		required := defaultExpr == "" && !pyTypeIsOptional(typeExpr)
		repeated := pyTypeIsRepeated(typeExpr)
		wireName := name
		if alias != "" {
			wireName = alias
		}
		shape.Fields = append(shape.Fields, ShapeField{
			Name:     wireName,
			Type:     typeExpr,
			JSONTag:  alias,
			Required: required,
			Repeated: repeated,
			Comment:  comment,
		})
	}
	if len(shape.Fields) == 0 {
		return nil
	}
	return shape
}

func pyTypeIsOptional(t string) bool {
	t = strings.TrimSpace(t)
	if strings.HasPrefix(t, "Optional[") {
		return true
	}
	// `str | None` / `None | str` — Python 3.10+ union syntax.
	for _, p := range strings.Split(t, "|") {
		if strings.TrimSpace(p) == "None" {
			return true
		}
	}
	return false
}

func pyTypeIsRepeated(t string) bool {
	t = strings.TrimSpace(t)
	// list[Foo], List[Foo], tuple[Foo, ...], Sequence[Foo]
	for _, pfx := range []string{"list[", "List[", "tuple[", "Tuple[", "Sequence[", "Iterable[", "set[", "Set[", "FrozenSet["} {
		if strings.HasPrefix(t, pfx) {
			return true
		}
	}
	return false
}

func leadingWhitespace(s string) string {
	for i, r := range s {
		if r != ' ' && r != '\t' {
			return s[:i]
		}
	}
	return s
}