diff options
| -rw-r--r-- | cmd/rescribe/main.go | 55 | ||||
| -rw-r--r-- | internal/pipeline/get.go | 58 | ||||
| -rw-r--r-- | local.go | 2 | 
3 files changed, 108 insertions, 7 deletions
| diff --git a/cmd/rescribe/main.go b/cmd/rescribe/main.go index e3781cb..c309367 100644 --- a/cmd/rescribe/main.go +++ b/cmd/rescribe/main.go @@ -10,6 +10,7 @@ package main  import (  	"flag"  	"fmt" +	"io/ioutil"  	"log"  	"os"  	"path/filepath" @@ -88,7 +89,7 @@ func main() {  	bookdir := flag.Arg(0)  	var bookname string -	if flag.NArg() > 2 { +	if flag.NArg() > 1 {  		bookname = flag.Arg(1)  	} else {  		bookname = filepath.Base(bookdir) @@ -102,22 +103,41 @@ func main() {  		verboselog = log.New(n, "", 0)  	} +	tempdir, err := ioutil.TempDir("", "bookpipeline") +	if err != nil { +		log.Fatalln("Error setting up temporary directory:", err) +	} +  	var conn Pipeliner -	// TODO: set tmpdir to a specific random thing for this run only -	conn = &bookpipeline.LocalConn{Logger: verboselog} +	conn = &bookpipeline.LocalConn{Logger: verboselog, TempDir: tempdir}  	conn.Log("Setting up session") -	err := conn.Init() +	err = conn.Init()  	if err != nil {  		log.Fatalln("Error setting up connection:", err)  	}  	conn.Log("Finished setting up session") -	uploadbook(bookdir, bookname, *training, conn) +	fmt.Printf("Copying book to pipeline\n") +	err = uploadbook(bookdir, bookname, *training, conn) +	if err != nil { +		log.Fatalln(err) +	} + +	fmt.Printf("Processing book (this may take some time)\n")  	processbook(*training, conn) -	// TODO: save book +	fmt.Printf("Saving finished book to %s\n", bookname) +	err = downloadbook(bookname, conn) +	if err != nil { +		log.Fatalln(err) +	} + +	err = os.RemoveAll(tempdir) +	if err != nil { +		log.Fatalf("Error removing temporary directory %s: %v", tempdir, err) +	}  }  func uploadbook(dir string, name string, training string, conn Pipeliner) error { @@ -142,6 +162,29 @@ func uploadbook(dir string, name string, training string, conn Pipeliner) error  	return nil  } +func downloadbook(name string, conn Pipeliner) error { +	err := os.MkdirAll(name, 0755) +	if err != nil { +		log.Fatalln("Failed to create directory", name, err) +	} + +	err = pipeline.DownloadBestPages(name, conn) +	if err != nil { +		return fmt.Errorf("Error downloading best pages: %v", err) +	} + +	err = pipeline.DownloadPdfs(name, conn) +	if err != nil { +		return fmt.Errorf("Error downloading PDFs: %v", err) +	} + +	err = pipeline.DownloadAnalyses(name, conn) +	if err != nil { +		return fmt.Errorf("Error downloading analyses: %v", err) +	} + +	return nil +}  func processbook(training string, conn Pipeliner) {  	origPattern := regexp.MustCompile(`[0-9]{4}.jpg$`) diff --git a/internal/pipeline/get.go b/internal/pipeline/get.go new file mode 100644 index 0000000..8492d99 --- /dev/null +++ b/internal/pipeline/get.go @@ -0,0 +1,58 @@ +// Copyright 2019 Nick White. +// Use of this source code is governed by the GPLv3 +// license that can be found in the LICENSE file. + +package pipeline + +import ( +	"bufio" +	"fmt" +	"os" +	"path/filepath" +) + +func DownloadBestPages(name string, conn Pipeliner) error { +	fn := filepath.Join(name, "best") +	err := conn.Download(conn.WIPStorageId(), fn, fn) +	if err != nil { +		return fmt.Errorf("Failed to download 'best' file: %v", err) +	} +	f, err := os.Open(fn) +	if err != nil { +		return fmt.Errorf("Failed to open best file: %v", err) +	} +	defer f.Close() + +	s := bufio.NewScanner(f) +	for s.Scan() { +		fn = filepath.Join(name, s.Text()) +		err = conn.Download(conn.WIPStorageId(), fn, fn) +		if err != nil { +			return fmt.Errorf("Failed to download file %s: %v", fn, err) +		} +	} + +	return nil +} + +func DownloadPdfs(name string, conn Pipeliner) error { +	for _, suffix := range []string{".colour.pdf", ".binarised.pdf"} { +		fn := filepath.Join(name, name+suffix) +		err := conn.Download(conn.WIPStorageId(), fn, fn) +		if err != nil { +			return fmt.Errorf("Failed to download PDF %s: %v", fn, err) +		} +	} +	return nil +} + +func DownloadAnalyses(name string, conn Pipeliner) error { +	for _, a := range []string{"conf", "graph.png"} { +		fn := filepath.Join(name, a) +		err := conn.Download(conn.WIPStorageId(), fn, fn) +		if err != nil { +			return fmt.Errorf("Failed to download analysis file %s: %v", fn, err) +		} +	} +	return nil +} @@ -36,7 +36,7 @@ func (a *LocalConn) MinimalInit() error {  	if a.TempDir == "" {  		a.TempDir = filepath.Join(os.TempDir(), "bookpipeline")  	} -	err = os.Mkdir(a.TempDir, 0700) +	err = os.MkdirAll(a.TempDir, 0700)  	if err != nil && !os.IsExist(err) {  		return fmt.Errorf("Error creating temporary directory: %v", err)  	} | 
