py2uast2pb.go

package main

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

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

var pythonMatch = "/usr/lib/python3.6/*.py"
var serverAddr = "0.0.0.0:9432"
var timeout time.Duration = 2

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

func connect(addr string, timeout time.Duration) (protocol.ProtocolServiceClient, *grpc.ClientConn, error) {
	conn, err := grpc.Dial(addr, grpc.WithTimeout(time.Second*timeout),
		grpc.WithInsecure())
	if err != nil {
		return nil, conn, err
	}

	client := protocol.NewProtocolServiceClient(conn)
	return client, conn, nil
}

func main() {
	client, conn, err := connect(serverAddr, timeout)
	defer conn.Close()
	if err != nil {
		fmt.Println("Could not connect to the server, good bye")
		os.Exit(1)
	}

	pyfiles, err := filepath.Glob(pythonMatch)
	if err != nil {
		fmt.Println("filepath.Glob failed for ", pythonMatch)
		os.Exit(0)
	}
	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",
		}

		// It will hang just here after the first FATAL
		resp, err := client.ParseUAST(context.TODO(), req)
		if err != nil {
			fmt.Println("ParseUAST failed for: ", fname)
			continue
		}

		if resp.Status.String() == "fatal" {
			fmt.Println("FATAL error with the file [", fname, "]:", resp.Errors)

			// Reconnect (also hangs with it commented)
			conn.Close()
			client, conn, err = connect(serverAddr, timeout)
			if err != nil {
				fmt.Println("Could not reconnect to server, sorry")
				os.Exit(1)
			}
			continue
		}

		if resp.UAST == nil {
			fmt.Println("Something failed with the file [", fname, "], UAST is nil")
			continue
		}

		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)
	}
}