fix(protogen): map Python dependency imports

.github/workflows/tests.yml

@@ -19,6 +19,7 @@ jobs:
       - uses: actions/setup-go@v7
         with:
           go-version: '1.26'
+      - uses: astral-sh/setup-uv@v7
       - run: bun install
       - name: Cache tools
         uses: actions/cache@v6

protogen/cache.go

@@ -97,11 +97,15 @@ func (c *Cache) SetToolVersions(versions string) {
 // - The file content hash has changed
 // - The protoc flags have changed
 // - Force is true
-func (c *Cache) NeedsRegeneration(packageKey string, protoFiles []string, projectDir string, flagsHash string, force bool) (bool, error) {
+func (c *Cache) NeedsRegeneration(packageKey string, protoFiles []string, projectDir string, flagsHash string, toolVersions string, force bool) (bool, error) {
 	if force {
 		return true, nil
 	}
 
+	if c.ToolVersions != toolVersions {
+		return true, nil
+	}
+
 	// Check if flags changed
 	if c.ProtocFlagsHash != flagsHash {
 		return true, nil

protogen/cache_test.go

@@ -1,9 +1,14 @@
 package protogen
 
 import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
 	"os"
 	"path/filepath"
 	"slices"
+	"strings"
 	"testing"
 )
 
@@ -194,3 +199,115 @@ func TestHashProtoFilesPortableAndContentSensitive(t *testing.T) {
 		t.Fatalf("content hash must change when a proto source changes")
 	}
 }
+
+func TestToolVersionChangeInvalidatesCachedPackage(t *testing.T) {
+	dir, files := writeProtoTree(t, map[string]string{
+		"foo.proto": "syntax = \"proto3\";\npackage foo;\n",
+	})
+	cache := NewCache()
+	cache.ProtocFlagsHash = "flags"
+	cache.ToolVersions = "starpc-python=old"
+	if err := cache.UpdatePackage("example/foo", files, nil, dir); err != nil {
+		t.Fatal(err)
+	}
+
+	stale, err := cache.NeedsRegeneration(
+		"example/foo", files, dir, "flags", "starpc-python=old", false,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stale {
+		t.Fatal("matching tool state invalidated cached output")
+	}
+	stale, err = cache.NeedsRegeneration(
+		"example/foo", files, dir, "flags", "starpc-python=new", false,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !stale {
+		t.Fatal("changed tool state reused cached output")
+	}
+}
+
+func TestGetToolVersionsIncludesUVLock(t *testing.T) {
+	dir := t.TempDir()
+	lock := []byte("version = 1\nrevision = 3\n")
+	if err := os.WriteFile(filepath.Join(dir, "uv.lock"), lock, 0o644); err != nil {
+		t.Fatal(err)
+	}
+	generator := &Generator{
+		ProjectDir: dir,
+		Config:     NewConfig(),
+		Plugins:    &Plugins{StarpcPython: &Plugin{}},
+	}
+	versions := generator.getToolVersions()
+	want := sha256.Sum256(lock)
+	if !strings.Contains(versions, "uv.lock="+hex.EncodeToString(want[:])) {
+		t.Fatalf("tool versions omit uv.lock digest: %q", versions)
+	}
+}
+
+func TestGetToolVersionsIgnoresUVLockWithoutStarpcPython(t *testing.T) {
+	dir := t.TempDir()
+	if err := os.WriteFile(filepath.Join(dir, "uv.lock"), []byte("changed"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	generator := &Generator{
+		ProjectDir: dir,
+		Config:     NewConfig(),
+		Plugins:    &Plugins{},
+	}
+	if versions := generator.getToolVersions(); strings.Contains(versions, "uv.lock=") {
+		t.Fatalf("unselected Python plugin invalidated tool state: %q", versions)
+	}
+}
+
+func TestGenerateFailureDoesNotPersistToolVersions(t *testing.T) {
+	dir := t.TempDir()
+	vendorDir := filepath.Join(dir, "vendor")
+	if err := os.MkdirAll(vendorDir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(dir, "bad.proto"), []byte("not protobuf"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(filepath.Join(dir, "uv.lock"), []byte("changed plugin"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	config := NewConfig()
+	config.ProjectDir = dir
+	config.Targets = []string{"bad.proto"}
+	cache := NewCache()
+	cache.ToolVersions = "accepted"
+	plugins := &Plugins{
+		Languages:    Languages{LanguagePython: {}},
+		RPCLibraries: RPCLibraries{RPCLibraryStarpcPython: {}},
+		StarpcPython: &Plugin{
+			Name:       "starpc-python",
+			BinaryName: "protoc-gen-starpc-python",
+			Path:       "/bin/false",
+			Type:       PluginTypePython,
+			OutFlag:    "starpc-python_out",
+		},
+	}
+	generator := &Generator{
+		Config:     config,
+		Plugins:    plugins,
+		Cache:      cache,
+		ProjectDir: dir,
+		ModuleDir:  dir,
+		ModulePath: "example.com/project",
+		VendorDir:  vendorDir,
+		OutDir:     vendorDir,
+		Stdout:     io.Discard,
+		Stderr:     io.Discard,
+	}
+	if err := generator.Generate(context.Background()); err == nil {
+		t.Fatal("invalid source generation unexpectedly succeeded")
+	}
+	if cache.ToolVersions != "accepted" {
+		t.Fatalf("failed generation persisted tool state %q", cache.ToolVersions)
+	}
+}

protogen/generator.go

@@ -4,6 +4,8 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"os"
@@ -125,9 +127,8 @@ func (g *Generator) Generate(ctx context.Context) error {
 		fmt.Fprintf(g.Stdout, "Found %d proto files\n", len(protoFiles))
 	}
 
-	// Get tool versions for cache invalidation
+	// Get tool versions for cache invalidation.
 	toolVersions := g.getToolVersions()
-	g.Cache.SetToolVersions(toolVersions)
 
 	// Build protoc arguments
 	protocArgs := g.buildProtocArgs()
@@ -157,7 +158,7 @@ func (g *Generator) Generate(ctx context.Context) error {
 		currentPackages[packageKey] = struct{}{}
 
 		// Check if regeneration is needed
-		needsRegen, err := g.Cache.NeedsRegeneration(packageKey, files, g.ProjectDir, flagsHash, g.Config.Force)
+		needsRegen, err := g.Cache.NeedsRegeneration(packageKey, files, g.ProjectDir, flagsHash, toolVersions, g.Config.Force)
 		if err != nil {
 			return fmt.Errorf("failed to check cache for %s: %w", dir, err)
 		}
@@ -247,6 +248,7 @@ func (g *Generator) Generate(ctx context.Context) error {
 	g.Cache.CleanOrphanedPackages(currentPackages)
 
 	g.Cache.SetProtocFlags(protocArgs, g.ModuleDir)
+	g.Cache.SetToolVersions(toolVersions)
 	// Save cache
 	cacheFile, _ := g.Config.GetCacheFilePath()
 	if err := g.Cache.Save(cacheFile); err != nil {
@@ -453,6 +455,13 @@ func (g *Generator) getToolVersions() string {
 		}
 	}
 
+	if g.Plugins != nil && g.Plugins.StarpcPython != nil {
+		if data, err := os.ReadFile(filepath.Join(g.ProjectDir, "uv.lock")); err == nil {
+			digest := sha256.Sum256(data)
+			versions = append(versions, "uv.lock="+hex.EncodeToString(digest[:]))
+		}
+	}
+
 	return strings.Join(versions, ",")
 }
 

protogen/postprocess.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"slices"
 	"strings"
 )
 
@@ -221,13 +222,14 @@ func (p *PostProcessor) ProcessGoFile(filePath string) error {
 	return nil
 }
 
-// ProcessPythonFile rewrites local canonical module imports to adjacent imports.
+// ProcessPythonFile rewrites canonical Go module imports to the module-relative
+// packages installed by the current project and its vendored dependencies.
 func (p *PostProcessor) ProcessPythonFile(filePath string) error {
 	data, err := os.ReadFile(filePath)
 	if err != nil {
 		return err
 	}
-	prefix := strings.ReplaceAll(strings.ReplaceAll(p.ModulePath, "/", "."), "-", "_") + "."
+	prefixes := p.pythonImportPrefixes()
 	lines := strings.Split(string(data), "\n")
 	modified := false
 	for i, line := range lines {
@@ -236,24 +238,60 @@ func (p *PostProcessor) ProcessPythonFile(filePath string) error {
 			continue
 		}
 		from := strings.TrimPrefix(trimmed, "from ")
-		if !strings.HasPrefix(from, prefix) {
-			continue
-		}
-		rest := strings.TrimPrefix(from, prefix)
-		space := strings.IndexByte(rest, ' ')
-		if space <= 0 {
-			continue
+		for _, prefix := range prefixes {
+			if !strings.HasPrefix(from, prefix) {
+				continue
+			}
+			rest := strings.TrimPrefix(from, prefix)
+			if strings.IndexByte(rest, ' ') <= 0 {
+				continue
+			}
+			indent := line[:len(line)-len(trimmed)]
+			lines[i] = indent + "from " + rest
+			modified = true
+			break
 		}
-		indent := line[:len(line)-len(trimmed)]
-		lines[i] = indent + "from " + rest
-		modified = true
 	}
 	if !modified {
 		return nil
 	}
 	return os.WriteFile(filePath, []byte(strings.Join(lines, "\n")), 0o644) //nolint:gosec
 }
 
+func (p *PostProcessor) pythonImportPrefixes() []string {
+	modules := map[string]struct{}{p.ModulePath: {}}
+	modulesFile := filepath.Join(p.VendorDir, "modules.txt")
+	if data, err := os.ReadFile(modulesFile); err == nil {
+		scanner := bufio.NewScanner(bytes.NewReader(data))
+		for scanner.Scan() {
+			line := scanner.Text()
+			if !strings.HasPrefix(line, "# ") || strings.HasPrefix(line, "## ") {
+				continue
+			}
+			fields := strings.Fields(strings.TrimPrefix(line, "# "))
+			if len(fields) < 2 || (fields[1] != "=>" && !strings.HasPrefix(fields[1], "v")) {
+				continue
+			}
+			modules[fields[0]] = struct{}{}
+		}
+	}
+	prefixes := make([]string, 0, len(modules))
+	for module := range modules {
+		if module == "" {
+			continue
+		}
+		prefix := strings.ReplaceAll(strings.ReplaceAll(module, "/", "."), "-", "_")
+		prefixes = append(prefixes, prefix+".")
+	}
+	slices.SortFunc(prefixes, func(a, b string) int {
+		if len(a) != len(b) {
+			return len(b) - len(a)
+		}
+		return strings.Compare(a, b)
+	})
+	return prefixes
+}
+
 // ProcessTsFile processes a TypeScript file.
 // Rewrites relative import paths to @go/ format for vendor dependencies.
 //

protogen/postprocess_test.go

@@ -122,3 +122,53 @@ func TestProcessPythonFileRewritesPyi(t *testing.T) {
 		t.Fatalf("got %q", got)
 	}
 }
+
+func TestProcessPythonFileRewritesVendoredModuleImports(t *testing.T) {
+	projectDir := t.TempDir()
+	vendorDir := filepath.Join(projectDir, "vendor")
+	if err := os.MkdirAll(vendorDir, 0o755); err != nil {
+		t.Fatal(err)
+	}
+	modules := `# github.com/aperturerobotics/starpc v0.52.0
+## explicit; go 1.25.0
+# github.com/aperturerobotics/starpc/extensions v0.1.0
+## explicit; go 1.25.0
+# github.com/foo/bar v1.0.0
+## explicit; go 1.25.0
+github.com/foo/bar/pkg
+`
+	if err := os.WriteFile(filepath.Join(vendorDir, "modules.txt"), []byte(modules), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	file := filepath.Join(projectDir, "resource_pb2.py")
+	content := `from github.com.s4wave.spacewave.local import local_pb2
+from github.com.aperturerobotics.starpc.rpcstream import rpcstream_pb2
+from github.com.aperturerobotics.starpc.extensions.echo import echo_pb2
+from github.com.foo.bar.pkg import package_pb2
+from github.com.foo.barista.foo import barista_pb2
+from google.protobuf import timestamp_pb2
+from undeclared.example import example_pb2
+`
+	if err := os.WriteFile(file, []byte(content), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	pp := NewPostProcessor(projectDir, vendorDir, "github.com/s4wave/spacewave", nil, false)
+	if err := pp.ProcessPythonFile(file); err != nil {
+		t.Fatal(err)
+	}
+	got, err := os.ReadFile(file)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := `from local import local_pb2
+from rpcstream import rpcstream_pb2
+from echo import echo_pb2
+from pkg import package_pb2
+from github.com.foo.barista.foo import barista_pb2
+from google.protobuf import timestamp_pb2
+from undeclared.example import example_pb2
+`
+	if string(got) != want {
+		t.Fatalf("unexpected dependency import rewrite:\n%s", got)
+	}
+}