Created
October 9, 2019 09:10
-
-
Save ovechkin-dm/b06f29efef0f4d8f50a6b73018127694 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"fmt" | |
"io/ioutil" | |
"os" | |
"os/exec" | |
"path/filepath" | |
) | |
const protoDir = "proto/" | |
var opts = []string{"-Iserving", "-Itensorflow"} | |
var cmds = []ProtocCmd{{ | |
PkgDir: protoDir + "tensorflow/core/example", | |
Inputs: []string{"tensorflow/tensorflow/core/example/*.proto"}, | |
}, { | |
PkgDir: protoDir + "tensorflow/core/framework", | |
Inputs: []string{"tensorflow/tensorflow/core/framework/*.proto"}, | |
}, { | |
PkgDir: protoDir + "tensorflow/core/lib/core", | |
Inputs: []string{"tensorflow/tensorflow/core/lib/core/*.proto"}, | |
}, { | |
GoOpts: "import_path=protobuf", | |
PkgDir: protoDir + "tensorflow/core/protobuf", | |
Inputs: []string{ | |
"tensorflow/tensorflow/core/protobuf/*.proto", | |
"tensorflow/tensorflow/stream_executor/*.proto", | |
}, | |
}, { | |
GoOpts: "plugins=grpc,import_path=serving", | |
PkgDir: protoDir + "tensorflow/serving", | |
Inputs: []string{ | |
"serving/tensorflow_serving/apis/*.proto", | |
"serving/tensorflow_serving/config/*.proto", | |
"serving/tensorflow_serving/core/*.proto", | |
"serving/tensorflow_serving/sources/storage_path/*.proto", | |
"serving/tensorflow_serving/util/*.proto", | |
}, | |
}} | |
func main() { | |
for _, cmd := range cmds { | |
fmt.Fprintln(os.Stderr, "==>", cmd.PkgDir) | |
if err := cmd.Run(); err != nil { | |
if e, ok := err.(*exec.ExitError); ok { | |
os.Exit(e.ExitCode()) | |
} | |
fmt.Fprintln(os.Stderr, err) | |
os.Exit(1) | |
} | |
} | |
} | |
// ProtocCmd executes protoc to generate sources for a single Go package. | |
type ProtocCmd struct { | |
GoOpts string // --go_out options | |
PkgDir string // Final output directory | |
Inputs []string // Input files or glob patterns | |
} | |
func (pc *ProtocCmd) Run() error { | |
// Use a temporary protoc output directory | |
root := filepath.Dir(pc.PkgDir) | |
os.MkdirAll(root, 0777) | |
tmp, err := ioutil.TempDir(root, filepath.Base(pc.PkgDir)+".") | |
if err != nil { | |
return err | |
} | |
defer os.RemoveAll(tmp) | |
// Run protoc | |
cmd := exec.Command("protoc", opts...) | |
cmd.Args = append(cmd.Args, "--go_out="+pc.GoOpts+":"+tmp) | |
for _, in := range pc.Inputs { | |
files, err := filepath.Glob(in) | |
if err != nil { | |
return err | |
} | |
cmd.Args = append(cmd.Args, files...) | |
} | |
cmd.Stdin, cmd.Stdout, cmd.Stderr = os.Stdin, os.Stdout, os.Stderr | |
err = cmd.Run() | |
// Move generated files to PkgDir | |
os.RemoveAll(pc.PkgDir) | |
if err := os.MkdirAll(pc.PkgDir, 0777); err != nil { | |
return err | |
} | |
walkErr := filepath.Walk(tmp, func(path string, fi os.FileInfo, err error) error { | |
if err == nil && fi.Mode().IsRegular() { | |
err = os.Rename(path, filepath.Join(pc.PkgDir, fi.Name())) | |
} | |
return err | |
}) | |
if err == nil { | |
err = walkErr | |
} | |
return err | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment