Generate Python dataclasses from Go status structs

Makefile

@@ -798,3 +798,8 @@ docs-%:
 ## docs-run: Build and serve the documentation
 ## docs-clean: Clean the docs build artifacts
 	cd docs && $(MAKE) -f Makefile.sp sp-$* ALLFILES='*.md **/*.md'
+
+.PHONY: getfields
+getfields:
+	go run ./cmd/getfields/ >structs.py
+	uvx [email protected] format --line-length=99 --config "format.quote-style = 'single'" structs.py

cmd/getfields/main.go

@@ -0,0 +1,387 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"log"
+	"maps"
+	"os"
+	"slices"
+	"strings"
+
+	"github.com/juju/juju/cmd/juju/status"
+)
+
+const maxLineLength = 99
+
+var classDocstrings = map[string]string{
+	"Status": `Parsed version of the status object returned by "juju status --format=json".`,
+}
+
+var additionalMethods = map[string]string{
+	"AppStatus": `
+    @property
+    def is_active(self) -> bool:
+        """Report whether the application status for this app is "active"."""
+        return self.app_status.current == 'active'
+
+    @property
+    def is_blocked(self) -> bool:
+        """Report whether the application status for this app is "blocked"."""
+        return self.app_status.current == 'blocked'
+
+    @property
+    def is_error(self) -> bool:
+        """Report whether the application status for this app is "error"."""
+        return self.app_status.current == 'error'
+
+    @property
+    def is_maintenance(self) -> bool:
+        """Report whether the application status for this app is "maintenance"."""
+        return self.app_status.current == 'maintenance'
+
+    @property
+    def is_waiting(self) -> bool:
+        """Report whether the application status for this app is "waiting"."""
+        return self.app_status.current == 'waiting'
+`,
+	"UnitStatus": `
+
+    @property
+    def is_active(self) -> bool:
+        """Report whether the workload status for this unit status is "active"."""
+        return self.workload_status.current == 'active'
+
+    @property
+    def is_blocked(self) -> bool:
+        """Report whether the workload status for this unit status is "blocked"."""
+        return self.workload_status.current == 'blocked'
+
+    @property
+    def is_error(self) -> bool:
+        """Report whether the workload status for this unit status is "error"."""
+        return self.workload_status.current == 'error'
+
+    @property
+    def is_maintenance(self) -> bool:
+        """Report whether the workload status for this unit status is "maintenance"."""
+        return self.workload_status.current == 'maintenance'
+
+    @property
+    def is_waiting(self) -> bool:
+        """Report whether the workload status for this unit status is "waiting"."""
+        return self.workload_status.current == 'waiting'
+`,
+}
+
+func main() {
+	structs := status.GetFields()
+
+	classes := make(map[string]string)
+	successors := make(map[string][]string)
+
+	structNames := slices.Sorted(maps.Keys(structs))
+
+	requireArgsStructs := make(map[string]bool)
+	for _, name := range structNames {
+		requireArgs := false
+		for _, field := range structs[name] {
+			if field.JSONField == "" {
+				continue
+			}
+			if !field.OmitEmpty {
+				requireArgs = true
+			}
+		}
+		className := getClassName(name)
+		if requireArgs {
+			requireArgsStructs[className] = requireArgs
+		}
+	}
+
+	var allNames []string
+	for _, name := range structNames {
+		var buf bytes.Buffer
+
+		fmt.Fprintln(&buf, "\n\[email protected](frozen=True, kw_only=True)")
+		className := getClassName(name)
+		allNames = append(allNames, className)
+		fmt.Fprintf(&buf, "class %s:\n", className)
+		if docstring, ok := classDocstrings[className]; ok {
+			fmt.Fprintf(&buf, `    """%s"""`+"\n", docstring)
+		}
+		successors[className] = nil
+		var required []string
+		var optional []string
+		for _, field := range structs[name] {
+			if field.JSONField == "" {
+				continue
+			}
+			pythonField := getPythonField(field.JSONField)
+			pythonType := ""
+			if _, ok := structs[field.Type]; ok {
+				pythonType, _ = getPythonType(structs, field.Type)
+				successors[className] = append(successors[className], pythonType)
+			} else {
+				var base string
+				pythonType, base = getPythonType(structs, field.Type)
+				if base != "" {
+					successors[className] = append(successors[className], base)
+				}
+			}
+			if field.OmitEmpty {
+				switch {
+				case strings.HasPrefix(pythonType, "list["):
+					pythonType += " = dataclasses.field(default_factory=list)"
+				case strings.HasPrefix(pythonType, "dict[str, "):
+					pythonType += " = dataclasses.field(default_factory=dict)"
+				case pythonType == "bool":
+					pythonType += " = False"
+				case pythonType == "str":
+					pythonType += " = ''"
+				case pythonType == "int":
+					pythonType += " = 0"
+				default:
+					if requireArgsStructs[pythonType] {
+						pythonType += " | None = None"
+					} else {
+						pythonType += fmt.Sprintf(" = dataclasses.field(default_factory=%s)", pythonType)
+					}
+				}
+			}
+			line := fmt.Sprintf("    %s: %s\n", pythonField, pythonType)
+			if field.OmitEmpty {
+				optional = append(optional, line)
+			} else {
+				required = append(required, line)
+			}
+		}
+		for _, line := range required {
+			fmt.Fprint(&buf, line)
+		}
+		if len(required) > 0 && len(optional) > 0 {
+			fmt.Fprintln(&buf)
+		}
+		for _, line := range optional {
+			fmt.Fprint(&buf, line)
+		}
+
+		fmt.Fprintf(&buf, "\n    @classmethod\n")
+		fmt.Fprintf(&buf, "    def _from_dict(cls, d: dict[str, Any]) -> %s:\n", className)
+
+		hasErr := false
+		for _, field := range structs[name] {
+			if field.JSONField == "" && field.Name == "Err" && field.Type == "error" {
+				hasErr = true
+			}
+		}
+		// UnitStatus is a special case, has Err as .WorkloadStatusInfo.Err
+		if hasErr || className == "UnitStatus" {
+			fmt.Fprintln(&buf, "        if 'status-error' in d:")
+			fmt.Fprintln(&buf, "            raise StatusError(d['status-error'])")
+		}
+
+		fmt.Fprintln(&buf, "        return cls(")
+		optional = optional[:0]
+		required = required[:0]
+		for _, field := range structs[name] {
+			if field.JSONField == "" {
+				if field.Name == "Err" && field.Type == "error" {
+					continue
+				}
+				fmt.Fprintf(os.Stderr, "skipped field: %s.%s %s\n", name, field.Name, field.Type)
+				continue
+			}
+			pythonField := getPythonField(field.JSONField)
+			pythonType, _ := getPythonType(structs, field.Type)
+			dictGetter := getDictGetter(pythonType, field.JSONField, field.OmitEmpty, requireArgsStructs[pythonType])
+			line := fmt.Sprintf("            %s=%s,", pythonField, dictGetter)
+			if len(line) > maxLineLength {
+				// So that Ruff formats these lines nicely
+				line = fmt.Sprintf("            %s=(%s),", pythonField, dictGetter)
+			}
+			if field.OmitEmpty {
+				optional = append(optional, line+"\n")
+			} else {
+				required = append(required, line+"\n")
+			}
+		}
+
+		for _, line := range required {
+			fmt.Fprint(&buf, line)
+		}
+		for _, line := range optional {
+			fmt.Fprint(&buf, line)
+		}
+
+		fmt.Fprintln(&buf, "        )")
+
+		if methods, ok := additionalMethods[className]; ok {
+			fmt.Fprint(&buf, methods)
+		}
+
+		classes[className] = buf.String()
+		buf.Reset()
+	}
+
+	var order []string
+	for _, names := range tarjanSort(successors) {
+		if len(names) > 1 {
+			panic(fmt.Sprintf("dependency loop: %s", strings.Join(names, ", ")))
+		}
+		order = append(order, names[0])
+	}
+
+	allNames = append(allNames, "StatusError")
+	slices.Sort(allNames)
+
+	fmt.Print(`"""Dataclasses used to hold parsed output from "juju status --format=json"."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+__all__ = [
+`)
+	for _, name := range allNames {
+		fmt.Printf("    \"%s\",\n", name)
+	}
+	fmt.Print(`
+]
+
+
+class StatusError(Exception):
+    """Raised when "juju status" returns a status-error for certain types."""
+`)
+	for _, name := range order {
+		fmt.Print(classes[name])
+	}
+
+	f, err := os.Create("structs.json")
+	if err != nil {
+		log.Fatal(err)
+	}
+	b, err := json.MarshalIndent(structs, "", "  ")
+	if err != nil {
+		log.Fatal(err)
+	}
+	_, err = f.Write(b)
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = f.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func getClassName(goType string) string {
+	className := strings.Title(goType)
+	className = strings.ReplaceAll(className, "Application", "App")
+	if className == "FormattedStatus" {
+		className = "Status"
+	}
+	return className
+}
+
+func getPythonField(s string) string {
+	s = strings.ReplaceAll(s, "-", "_")
+	s = strings.ReplaceAll(s, "application", "app")
+	s = strings.ToLower(s)
+	return s
+}
+
+func getPythonType(structs map[string][]status.FieldInfo, goType string) (string, string) {
+	switch {
+	case strings.HasPrefix(goType, "[]"):
+		inner, base := getPythonType(structs, goType[2:])
+		return "list[" + inner + "]", base
+	case strings.HasPrefix(goType, "map[string]"):
+		inner, base := getPythonType(structs, goType[11:])
+		return "dict[str, " + inner + "]", base
+	case goType == "string":
+		return "str", ""
+	case goType == "bool":
+		return "bool", ""
+	case goType == "int" || goType == "uint64":
+		return "int", ""
+	default:
+		if _, ok := structs[goType]; !ok {
+			fmt.Fprintf(os.Stderr, "# unhandled Go type: %s\n", goType)
+		}
+		pythonType := getClassName(goType)
+		return pythonType, pythonType
+	}
+}
+
+func getDictGetter(pythonType string, jsonField string, omitEmpty bool, requireArgs bool) string {
+	s := fmt.Sprintf("d['%s']", jsonField)
+	orig := s
+	s = doType(pythonType, s)
+	if omitEmpty {
+		if s == orig {
+			// shortcut for simple value lookup
+			switch {
+			case strings.HasPrefix(pythonType, "list["):
+				return fmt.Sprintf("d.get('%s') or []", jsonField)
+			case strings.HasPrefix(pythonType, "dict[str, "):
+				return fmt.Sprintf("d.get('%s') or {}", jsonField)
+			case pythonType == "bool":
+				return fmt.Sprintf("d.get('%s') or False", jsonField)
+			case pythonType == "str":
+				return fmt.Sprintf("d.get('%s') or ''", jsonField)
+			case pythonType == "int":
+				return fmt.Sprintf("d.get('%s') or 0", jsonField)
+			default:
+				if !requireArgs {
+					return fmt.Sprintf("d.get('%s') or %s()", jsonField, pythonType)
+				}
+				return fmt.Sprintf("d.get('%s')", jsonField)
+			}
+		}
+		switch {
+		case strings.HasPrefix(pythonType, "list["):
+			s += fmt.Sprintf(" if '%s' in d else []", jsonField)
+		case strings.HasPrefix(pythonType, "dict[str, "):
+			s += fmt.Sprintf(" if '%s' in d else {}", jsonField)
+		case pythonType == "bool":
+			s += fmt.Sprintf(" if '%s' in d else False", jsonField)
+		case pythonType == "str":
+			s += fmt.Sprintf(" if '%s' in d else ''", jsonField)
+		case pythonType == "int":
+			s += fmt.Sprintf(" if '%s' in d else 0", jsonField)
+		default:
+			if !requireArgs {
+				s += fmt.Sprintf(" if '%s' in d else %s()", jsonField, pythonType)
+			} else {
+				s += fmt.Sprintf(" if '%s' in d else None", jsonField)
+			}
+		}
+	}
+	return s
+}
+
+func doType(pythonType string, value string) string {
+	switch {
+	case strings.HasPrefix(pythonType, "list["):
+		t := pythonType[5 : len(pythonType)-1]
+		inner := doType(t, "x")
+		if inner == "x" {
+			return value
+		}
+		return fmt.Sprintf("[%s for x in %s]", inner, value)
+	case strings.HasPrefix(pythonType, "dict[str, "):
+		t := pythonType[10 : len(pythonType)-1]
+		inner := doType(t, "v")
+		if inner == "v" {
+			return value
+		}
+		return fmt.Sprintf("{k: %s for k, v in %s.items()}", inner, value)
+	case pythonType == "str" || pythonType == "int" || pythonType == "bool":
+		return value
+	default:
+		return fmt.Sprintf("%s._from_dict(%s)", pythonType, value)
+	}
+}