py2ast2pb.go

package main

import (
    "context"
	"fmt"
    "os"
    "time"
	"path/filepath"
	"io/ioutil"

    "github.com/bblfsh/sdk/protocol"
    "google.golang.org/grpc"
)

var pythonMatch = "/usr/lib/python3.6/*.py"

// Note: start the server with:
// docker run --privileged -p 9432:9432 --name bblfsh bblfsh/server

func main() {
    // Connect to the running server
    conn, err := grpc.Dial("0.0.0.0:9432", grpc.WithTimeout(time.Second*2),
        grpc.WithInsecure())
    if err != nil {
        os.Exit(1)
    }
    client := protocol.NewProtocolServiceClient(conn)

	pyfiles, err := filepath.Glob(pythonMatch)
	if err != nil {
		fmt.Println("filepath.Glob failed for ", pythonMatch)
		return
	}
	for _, fname := range pyfiles {
		fmt.Println(fname)

		content, err := ioutil.ReadFile(fname)
		if err != nil {
			fmt.Println("Error: could not read file: ", fname)
			continue
		}
		contentstr := string(content)

		req := &protocol.ParseUASTRequest{Filename: fname,
			Content:  contentstr,
			Language: "python"}

		resp, err := client.ParseUAST(context.TODO(), req)
		if err != nil {
			fmt.Println("ParseUAST failed for: ", fname)
		}

		// FIXME: Crashes sometimes here, check resp.UAST? :(
		data, err := resp.UAST.Marshal()
		if err != nil {
			fmt.Println("Error: failed to export file: ", fname)
			continue
		}

		_, f := filepath.Split(fname)
		err = ioutil.WriteFile(f + ".pb", data, 00644)
	}
}