sift/cmd/sift/extractcmd.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
)

// maxExtractUpload mirrors the server's /v1/extract cap (a file over 100 MiB is
// refused on upload with 413), so reject it before uploading.
const maxExtractUpload = 100 << 20

type extractOpts struct {
	File, Pages, Grep, Question, Mode string
	JSON                              bool
}

const extractUsage = "usage: sift extract <file.pdf> [--pages a-b] [--grep RE] [--question Q] [--mode clean|digest] [--json]"

// parseExtractOpts parses `sift extract` args into extractOpts. It knows which
// flags take a value; the PDF path is the one positional.
func parseExtractOpts(args []string) (extractOpts, error) {
	flagArgs, pos, err := splitAroundPositional(args, extractValueFlags)
	if err != nil {
		return extractOpts{}, err
	}
	fs := flag.NewFlagSet("extract", flag.ContinueOnError)
	fs.SetOutput(io.Discard)
	modeF := fs.String("mode", "", "clean|digest (default digest if --question)")
	pagesF := fs.String("pages", "", "only these pages of a PDF (a-b or a)")
	grepF := fs.String("grep", "", "only lines matching RE in a PDF")
	questionF := fs.String("question", "", "focus the digest")
	asJSONF := fs.Bool("json", false, "print full /v1/extract envelope")
	if err := fs.Parse(flagArgs); err != nil {
		return extractOpts{}, err
	}
	if fs.NArg() != 0 {
		return extractOpts{}, fmt.Errorf("%s", extractUsage)
	}
	if len(pos) != 1 {
		return extractOpts{}, fmt.Errorf("%s", extractUsage)
	}
	return extractOpts{
		File:     pos[0],
		Mode:     *modeF,
		Pages:    *pagesF,
		Grep:     *grepF,
		Question: *questionF,
		JSON:     *asJSONF,
	}, nil
}

var extractValueFlags = map[string]bool{
	"mode": true, "pages": true, "grep": true, "question": true,
}

// extractURL builds the /v1/extract URL from base plus only the options that are
// set; mode defaults to digest when a question is given and no mode was.
func extractURL(base string, o extractOpts) string {
	v := url.Values{}
	if o.Pages != "" {
		v.Set("pages", o.Pages)
	}
	if o.Grep != "" {
		v.Set("grep", o.Grep)
	}
	if o.Question != "" {
		v.Set("question", o.Question)
	}
	mode := o.Mode
	if mode == "" && o.Question != "" {
		mode = "digest"
	}
	if mode != "" {
		v.Set("mode", mode)
	}
	u, err := url.Parse(base)
	if err != nil {
		return base + "/v1/extract?" + v.Encode()
	}
	u.Path = "/v1/extract"
	u.RawQuery = v.Encode()
	return u.String()
}

func extractMain() {
	o, err := parseExtractOpts(os.Args[1:])
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
	base, key, err := loadClient()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	f, err := os.Open(o.File)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	defer f.Close()
	st, err := f.Stat()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	if st.Size() > maxExtractUpload {
		fmt.Fprintf(os.Stderr, "extract: %s is %d bytes, over the %d MB limit\n", o.File, st.Size(), maxExtractUpload>>20)
		os.Exit(3)
	}
	req, err := http.NewRequest(http.MethodPost, extractURL(base, o), f)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	req.ContentLength = st.Size()
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Content-Type", "application/pdf")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	defer resp.Body.Close()
	var env map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&env); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	if resp.StatusCode != 200 {
		msg, _ := env["error"].(string)
		hint, _ := env["hint"].(string)
		if hint != "" {
			fmt.Fprintf(os.Stderr, "%s\nhint: %s\n", msg, hint)
		} else {
			fmt.Fprintln(os.Stderr, msg)
		}
		os.Exit(exitFor(resp.StatusCode, msg))
	}
	if o.JSON {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		_ = enc.Encode(env)
		return
	}
	if c, ok := env["content"].(string); ok {
		fmt.Print(c)
		if !strings.HasSuffix(c, "\n") {
			fmt.Println()
		}
	}
}