diff --git a/thrift/annotation.go b/thrift/annotation.go index 46221923..4c84271c 100644 --- a/thrift/annotation.go +++ b/thrift/annotation.go @@ -440,6 +440,17 @@ func extractNameSpaceToAnnos(ast *parser.Thrift) parser.Annotation { return ret } +// FilenameAnnotationKey is used for Option.PutThriftFilenameToAnnotation +const FilenameAnnotationKey = "thrift.filename" + +func extractThriftFilePathToAnnos(ast *parser.Thrift) parser.Annotation { + ret := parser.Annotation{ + Key: FilenameAnnotationKey, + } + ret.Values = append(ret.Values, ast.GetFilename()) + return ret +} + // injectAnnotation injects next annotation by appending. // NOTICE: the next annotation will be appended to the end of the current annotation. func injectAnnotations(origin *[]*parser.Annotation, next []parser.Annotation) error { diff --git a/thrift/idl.go b/thrift/idl.go index d406f141..b6e5160a 100644 --- a/thrift/idl.go +++ b/thrift/idl.go @@ -83,6 +83,14 @@ type Options struct { // `namespace go base` will got ["go", "base"] // NOTICE: at present, only StructDescriptor.Annotations() can get this PutNameSpaceToAnnotation bool + + // PutThriftFilenameToAnnotation indicates to extract the filename of one type + // and put it on the type's annotation. The annotion format is: + // - Key: "thrift.filename" (== FilenameAnnotationKey) + // - Values: pairs of Language and Name. for example: + // `// path := /a/b/c.thrift` will got ["/a/b/c.thrift"] + // NOTICE: at present, only StructDescriptor.Annotations() can get this + PutThriftFilenameToAnnotation bool } // NewDefaultOptions creates a default Options. @@ -559,6 +567,10 @@ func parseType(ctx context.Context, t *parser.Type, tree *parser.Thrift, cache c oannos = append(oannos, extractNameSpaceToAnnos(tree)) } + if opts.PutThriftFilenameToAnnotation { + oannos = append(oannos, extractThriftFilePathToAnnos(tree)) + } + // inject previous annotations injectAnnotations((*[]*parser.Annotation)(&st.Annotations), nextAnns) diff --git a/thrift/idl_test.go b/thrift/idl_test.go index 9066f3ac..2d9d90ac 100644 --- a/thrift/idl_test.go +++ b/thrift/idl_test.go @@ -20,6 +20,7 @@ import ( "context" "fmt" "math" + "path/filepath" "testing" "github.com/cloudwego/thriftgo/parser" @@ -303,6 +304,25 @@ func TestOptionPutNameSpaceToAnnotation(t *testing.T) { require.Equal(t, ns.Values, []string{"py", "py.base", "go", "go.base"}) } +func TestOptionPutThriftFilenameToAnnotation(t *testing.T) { + path := filepath.Join("..", "testdata", "idl", "example.thrift") + opt := Options{PutThriftFilenameToAnnotation: true} + descriptor, err := opt.NewDescritorFromPath(context.Background(), path) + require.NoError(t, err) + method, err := descriptor.LookupFunctionByMethod("ExampleMethod") + require.NoError(t, err) + req := method.Request().Struct().Fields()[0].Type() + annos := req.Struct().Annotations() + var filename *parser.Annotation + for i, a := range annos { + if a.Key == FilenameAnnotationKey { + filename = &annos[i] + break + } + } + require.Equal(t, filename.Values, []string{path}) +} + func TestNewFunctionDescriptorFromContent_absPath(t *testing.T) { content := ` include "/a/b/main.thrift"