diff --git a/protoc-gen-twirpy/generator/generator.go b/protoc-gen-twirpy/generator/generator.go index ccf1153..c8f6663 100644 --- a/protoc-gen-twirpy/generator/generator.go +++ b/protoc-gen-twirpy/generator/generator.go @@ -71,7 +71,7 @@ func buildTwirpServiceDescription(messagesToFiles map[string]string, fd *descrip FileName: name, } - imports := newImportBuilder(messagesToFiles) + imports := newImportBuilder(messagesToFiles, name) for _, service := range fd.GetService() { serviceURL := fmt.Sprintf("%s.%s", fd.GetPackage(), service.GetName()) diff --git a/protoc-gen-twirpy/generator/generator_test.go b/protoc-gen-twirpy/generator/generator_test.go index dac96ab..fb39493 100644 --- a/protoc-gen-twirpy/generator/generator_test.go +++ b/protoc-gen-twirpy/generator/generator_test.go @@ -9,13 +9,14 @@ import ( ) func TestImportBuilder(t *testing.T) { + // The current file being generated is in twirp/twitch/example/ testImportBuilder := newImportBuilder(map[string]string{ "twirp.twitch.example.Hat": "twirp/twitch/example/haberdasher.proto", "twirp.twitch.example.Price": "twirp/twitch/example/haberdasher.proto", "twirp.twitch.example.Color": "twirp/twitch/example_folder/haberdasher.proto", "twirp.twitch.example.Size": "twirp/twitch/example/haberdasher_extension.proto", "google.protobuf.Empty": "google/protobuf/empty.proto", - }) + }, "twirp/twitch/example/service.proto") testCases := []struct { typeToImport string @@ -24,26 +25,29 @@ func TestImportBuilder(t *testing.T) { expectedImport *TwirpImport }{ { + // Same package → relative import typeToImport: "twirp.twitch.example.Hat", importKey: "twirp.twitch.example.haberdasher_pb2", qualifiedImport: "_haberdasher_pb2.Hat", expectedImport: &TwirpImport{ - From: "twirp.twitch.example", + From: ".", Import: "haberdasher_pb2", Alias: "_haberdasher_pb2", }, }, { + // Same package, same file → reuses existing import typeToImport: "twirp.twitch.example.Price", importKey: "twirp.twitch.example.haberdasher_pb2", qualifiedImport: "_haberdasher_pb2.Price", expectedImport: &TwirpImport{ - From: "twirp.twitch.example", + From: ".", Import: "haberdasher_pb2", Alias: "_haberdasher_pb2", }, }, { + // Different package → absolute import typeToImport: "twirp.twitch.example.Color", importKey: "twirp.twitch.example_folder.haberdasher_pb2", qualifiedImport: "_haberdasher_pb2_1.Color", @@ -54,16 +58,18 @@ func TestImportBuilder(t *testing.T) { }, }, { + // Same package, different file → relative import typeToImport: "twirp.twitch.example.Size", importKey: "twirp.twitch.example.haberdasher_extension_pb2", qualifiedImport: "_haberdasher_extension_pb2.Size", expectedImport: &TwirpImport{ - From: "twirp.twitch.example", + From: ".", Import: "haberdasher_extension_pb2", Alias: "_haberdasher_extension_pb2", }, }, { + // External package → absolute import typeToImport: "google.protobuf.Empty", importKey: "google.protobuf.empty_pb2", qualifiedImport: "_empty_pb2.Empty", @@ -84,3 +90,39 @@ func TestImportBuilder(t *testing.T) { }) } } + +func TestImportBuilderRootLevel(t *testing.T) { + // currentFileName has no directory component: currentFileDir must be "" + // and root-level dependencies must produce absolute imports, not "." + ib := newImportBuilder(map[string]string{ + "mypackage.Foo": "types.proto", + }, "service.proto") + + qualified, err := ib.addImportAndQualify("mypackage.Foo") + assert.NoError(t, err) + assert.Equal(t, "_types_pb2.Foo", qualified) + // From must be "" (absolute import), not "." (relative) + assert.Equal(t, &TwirpImport{ + From: "", + Import: "types_pb2", + Alias: "_types_pb2", + }, ib.imports["types_pb2"]) +} + +func TestImportBuilderHyphenatedDirectory(t *testing.T) { + // Hyphens in directory names are normalised to underscores so that the + // directory string matches the Python module path produced by getModuleName. + ib := newImportBuilder(map[string]string{ + "my_org.my_service.Request": "my-org/my-service/types.proto", + }, "my-org/my-service/service.proto") + + qualified, err := ib.addImportAndQualify("my_org.my_service.Request") + assert.NoError(t, err) + assert.Equal(t, "_types_pb2.Request", qualified) + // Same package after hyphen normalisation → relative import + assert.Equal(t, &TwirpImport{ + From: ".", + Import: "types_pb2", + Alias: "_types_pb2", + }, ib.imports["my_org.my_service.types_pb2"]) +} diff --git a/protoc-gen-twirpy/generator/import_builder.go b/protoc-gen-twirpy/generator/import_builder.go index 714b0da..ae7e29c 100644 --- a/protoc-gen-twirpy/generator/import_builder.go +++ b/protoc-gen-twirpy/generator/import_builder.go @@ -9,13 +9,25 @@ type importBuilder struct { seenAliases map[string]struct{} imports map[string]*TwirpImport messagesToFile map[string]string + currentFileDir string } -func newImportBuilder(messagesToFile map[string]string) *importBuilder { +func newImportBuilder(messagesToFile map[string]string, currentFileName string) *importBuilder { + // Extract directory from proto filename: "crypto/v1.proto" -> "crypto" + dir := "" + parts := strings.Split(currentFileName, "/") + if len(parts) > 1 { + dirParts := parts[:len(parts)-1] + for i, p := range dirParts { + dirParts[i] = strings.ReplaceAll(p, "-", "_") + } + dir = strings.Join(dirParts, ".") + } return &importBuilder{ messagesToFile: messagesToFile, seenAliases: make(map[string]struct{}), imports: make(map[string]*TwirpImport), + currentFileDir: dir, } } @@ -34,6 +46,12 @@ func (ib *importBuilder) addImportAndQualify(typeToImport string) (string, error moduleNameSlice := strings.Split(moduleName, ".") strippedModuleName := moduleNameSlice[len(moduleNameSlice)-1] modulePath := strings.Join(moduleNameSlice[:len(moduleNameSlice)-1], ".") + // Use relative import when the module is in the same package as the file being generated. + // Require modulePath to be non-empty: an empty modulePath (root-level module) would + // otherwise match an empty currentFileDir and incorrectly emit a "." relative import. + if modulePath != "" && modulePath == ib.currentFileDir { + modulePath = "." + } alias := ib.generateAlias(strippedModuleName) ib.imports[moduleName] = &TwirpImport{ From: modulePath,