Repository: jhump/protoreflect
Pull request #333 (2020-07-06 19:41:52)

Add LookupImportProto
All that is used from `LookupImport` is the `FileDescriptorProto`, and its a lot of extra work to create a `desc.FileDescriptor` for `LookupImport` just so the parser can end up taking the `FileDescriptorProto` out of it


Patch:
diff --git a/desc/protoparse/errors.go b/desc/protoparse/errors.go
index 79e2adc..77ead4d 100644
--- a/desc/protoparse/errors.go
+++ b/desc/protoparse/errors.go
@@ -16,6 +16,9 @@ var ErrInvalidSource = errors.New("parse failed: invalid proto source")
 // indicates the file that had no syntax statement.
 var ErrNoSyntax = errors.New("no syntax specified; defaulting to proto2 syntax")
 
+// ErrLookupImportAndProtoSet is the error returned if both LookupImport and LookupImportProto are set.
+var ErrLookupImportAndProtoSet = errors.New("both LookupImport and LookupImportProto set")
+
 // ErrorReporter is responsible for reporting the given error. If the reporter
 // returns a non-nil error, parsing/linking will abort with that error. If the
 // reporter returns nil, parsing will continue, allowing the parser to try to
diff --git a/desc/protoparse/parser.go b/desc/protoparse/parser.go
index 6025bd0..ef27aff 100644
--- a/desc/protoparse/parser.go
+++ b/desc/protoparse/parser.go
@@ -118,6 +118,12 @@ type Parser struct {
 	// desc.LoadFileDescriptor.
 	LookupImport func(string) (*desc.FileDescriptor, error)
 
+	// LookupImportProto has the same functionality as LookupImport, however it returns
+	// a FileDescriptorProto instead of a FileDescriptor.
+	//
+	// It is an error to set both LookupImport and LookupImportProto.
+	LookupImportProto func(string) (*dpb.FileDescriptorProto, error)
+
 	// Used to create a reader for a given filename, when loading proto source
 	// file contents. If unset, os.Open is used. If ImportPaths is also empty
 	// then relative paths are will be relative to the process's current working
@@ -200,11 +206,15 @@ func (p Parser) ParseFiles(filenames ...string) ([]*desc.FileDescriptor, error)
 			return nil, ret
 		}
 	}
+	lookupImport, err := p.getLookupImport()
+	if err != nil {
+		return nil, err
+	}
 
 	protos := map[string]*parseResult{}
 	results := &parseResults{resultsByFilename: protos}
 	errs := newErrorHandler(p.ErrorReporter, p.WarningReporter)
-	parseProtoFiles(accessor, filenames, errs, true, true, results, p.LookupImport)
+	parseProtoFiles(accessor, filenames, errs, true, true, results, lookupImport)
 	if err := errs.getError(); err != nil {
 		return nil, err
 	}
@@ -269,10 +279,14 @@ func (p Parser) ParseFilesButDoNotLink(filenames ...string) ([]*dpb.FileDescript
 			return os.Open(name)
 		}
 	}
+	lookupImport, err := p.getLookupImport()
+	if err != nil {
+		return nil, err
+	}
 
 	protos := map[string]*parseResult{}
 	errs := newErrorHandler(p.ErrorReporter, p.WarningReporter)
-	parseProtoFiles(accessor, filenames, errs, false, p.ValidateUnlinkedFiles, &parseResults{resultsByFilename: protos}, p.LookupImport)
+	parseProtoFiles(accessor, filenames, errs, false, p.ValidateUnlinkedFiles, &parseResults{resultsByFilename: protos}, lookupImport)
 	if err := errs.getError(); err != nil {
 		return nil, err
 	}
@@ -301,6 +315,25 @@ func (p Parser) ParseFilesButDoNotLink(filenames ...string) ([]*dpb.FileDescript
 	return fds, nil
 }
 
+func (p Parser) getLookupImport() (func(string) (*dpb.FileDescriptorProto, error), error) {
+	if p.LookupImport != nil && p.LookupImportProto != nil {
+		return nil, ErrLookupImportAndProtoSet
+	}
+	if p.LookupImportProto != nil {
+		return p.LookupImportProto, nil
+	}
+	if p.LookupImport != nil {
+		return func(path string) (*dpb.FileDescriptorProto, error) {
+			value, err := p.LookupImport(path)
+			if value != nil {
+				return value.AsFileDescriptorProto(), err
+			}
+			return nil, err
+		}, nil
+	}
+	return nil, nil
+}
+
 func fixupFilenames(protos map[string]*parseResult) map[string]*parseResult {
 	// In the event that the given filenames (keys in the supplied map) do not
 	// match the actual paths used in 'import' statements in the files, we try
@@ -409,7 +442,7 @@ func fixupFilenames(protos map[string]*parseResult) map[string]*parseResult {
 	return revisedProtos
 }
 
-func parseProtoFiles(acc FileAccessor, filenames []string, errs *errorHandler, recursive, validate bool, parsed *parseResults, lookupImport func(string) (*desc.FileDescriptor, error)) {
+func parseProtoFiles(acc FileAccessor, filenames []string, errs *errorHandler, recursive, validate bool, parsed *parseResults, lookupImport func(string) (*dpb.FileDescriptorProto, error)) {
 	for _, name := range filenames {
 		parseProtoFile(acc, name, nil, errs, recursive, validate, parsed, lookupImport)
 		if errs.err != nil {
@@ -418,12 +451,12 @@ func parseProtoFiles(acc FileAccessor, filenames []string, errs *errorHandler, r
 	}
 }
 
-func parseProtoFile(acc FileAccessor, filename string, importLoc *SourcePos, errs *errorHandler, recursive, validate bool, parsed *parseResults, lookupImport func(string) (*desc.FileDescriptor, error)) {
+func parseProtoFile(acc FileAccessor, filename string, importLoc *SourcePos, errs *errorHandler, recursive, validate bool, parsed *parseResults, lookupImport func(string) (*dpb.FileDescriptorProto, error)) {
 	if parsed.has(filename) {
 		return
 	}
 	if lookupImport == nil {
-		lookupImport = func(string) (*desc.FileDescriptor, error) {
+		lookupImport = func(string) (*dpb.FileDescriptorProto, error) {
 			return nil, errors.New("no import lookup function")
 		}
 	}
@@ -442,7 +475,7 @@ func parseProtoFile(acc FileAccessor, filename string, importLoc *SourcePos, err
 	} else if d, lookupErr := lookupImport(filename); lookupErr == nil {
 		// This is a user-provided descriptor, which is acting similarly to a
 		// well-known import.
-		result = &parseResult{fd: proto.Clone(d.AsFileDescriptorProto()).(*dpb.FileDescriptorProto)}
+		result = &parseResult{fd: proto.Clone(d).(*dpb.FileDescriptorProto)}
 	} else if d, ok := standardImports[filename]; ok {
 		// it's a well-known import
 		// (we clone it to make sure we're not sharing state with other


Test patch:
diff --git a/desc/protoparse/parser_test.go b/desc/protoparse/parser_test.go
index bbf6e77..ccc351b 100644
--- a/desc/protoparse/parser_test.go
+++ b/desc/protoparse/parser_test.go
@@ -288,6 +288,25 @@ func TestParseFilesWithDependencies(t *testing.T) {
 			t.Errorf("Could not parse with a non-well-known import: %v", err)
 		}
 	})
+	t.Run("DependencyIncludedProto", func(t *testing.T) {
+		// Create a dependency-aware parser.
+		parser := Parser{
+			Accessor: FileContentsFromMap(contents),
+			LookupImportProto: func(imp string) (*dpb.FileDescriptorProto, error) {
+				if imp == "desc_test_wellknowntypes.proto" {
+					fileDescriptor, err := desc.LoadFileDescriptor(imp)
+					if err != nil {
+						return nil, err
+					}
+					return fileDescriptor.AsFileDescriptorProto(), nil
+				}
+				return nil, errors.New("unexpected filename")
+			},
+		}
+		if _, err := parser.ParseFiles("test.proto"); err != nil {
+			t.Errorf("Could not parse with a non-well-known import: %v", err)
+		}
+	})
 
 	// Establish that we *can not* parse the source file with a parser that
 	// did not register the dependency.