extract.go raw

   1  // Package extract provides text extraction from arbitrary file formats.
   2  // It shells out to external tools (pandoc, pdftotext) for complex formats
   3  // and reads plain text directly. Missing tools cause graceful fallback
   4  // to raw text or skip.
   5  package extract
   6  
   7  import (
   8  	"fmt"
   9  	"io"
  10  	"os"
  11  	"os/exec"
  12  	"path/filepath"
  13  	"strings"
  14  )
  15  
  16  // Extractor can extract plain text from a file.
  17  type Extractor interface {
  18  	// CanExtract reports whether this extractor handles the given path.
  19  	CanExtract(path string) bool
  20  
  21  	// Extract returns a reader of plain text content from the file.
  22  	// The caller must close the returned ReadCloser.
  23  	Extract(path string) (io.ReadCloser, error)
  24  }
  25  
  26  // Registry holds extractors in priority order.
  27  type Registry struct {
  28  	extractors []Extractor
  29  }
  30  
  31  // NewRegistry creates a registry with the standard set of extractors.
  32  func NewRegistry() *Registry {
  33  	return &Registry{
  34  		extractors: []Extractor{
  35  			PlainText{},
  36  			PDFExtractor{},
  37  			PandocExtractor{},
  38  		},
  39  	}
  40  }
  41  
  42  // Extract finds the first matching extractor and returns the text content.
  43  func (r *Registry) Extract(path string) (io.ReadCloser, error) {
  44  	for _, e := range r.extractors {
  45  		if e.CanExtract(path) {
  46  			return e.Extract(path)
  47  		}
  48  	}
  49  	return nil, fmt.Errorf("extract: no extractor for %s", path)
  50  }
  51  
  52  // PlainText reads .txt and .md files directly.
  53  type PlainText struct{}
  54  
  55  func (PlainText) CanExtract(path string) bool {
  56  	ext := strings.ToLower(filepath.Ext(path))
  57  	return ext == ".txt" || ext == ".md" || ext == ".text" || ext == ""
  58  }
  59  
  60  func (PlainText) Extract(path string) (io.ReadCloser, error) {
  61  	return os.Open(path)
  62  }
  63  
  64  // PDFExtractor uses pdftotext to extract text from PDFs.
  65  type PDFExtractor struct{}
  66  
  67  func (PDFExtractor) CanExtract(path string) bool {
  68  	return strings.ToLower(filepath.Ext(path)) == ".pdf"
  69  }
  70  
  71  func (PDFExtractor) Extract(path string) (io.ReadCloser, error) {
  72  	if _, err := exec.LookPath("pdftotext"); err != nil {
  73  		return nil, fmt.Errorf("extract: pdftotext not found: %w", err)
  74  	}
  75  	cmd := exec.Command("pdftotext", "-enc", "UTF-8", path, "-")
  76  	stdout, err := cmd.StdoutPipe()
  77  	if err != nil {
  78  		return nil, err
  79  	}
  80  	if err := cmd.Start(); err != nil {
  81  		return nil, err
  82  	}
  83  	return &cmdReadCloser{ReadCloser: stdout, cmd: cmd}, nil
  84  }
  85  
  86  // PandocExtractor uses pandoc to convert various formats to plain text.
  87  type PandocExtractor struct{}
  88  
  89  var pandocExts = map[string]bool{
  90  	".html":  true,
  91  	".htm":   true,
  92  	".docx":  true,
  93  	".epub":  true,
  94  	".rtf":   true,
  95  	".odt":   true,
  96  	".rst":   true,
  97  	".latex": true,
  98  	".tex":   true,
  99  	".org":   true,
 100  }
 101  
 102  func (PandocExtractor) CanExtract(path string) bool {
 103  	ext := strings.ToLower(filepath.Ext(path))
 104  	return pandocExts[ext]
 105  }
 106  
 107  func (PandocExtractor) Extract(path string) (io.ReadCloser, error) {
 108  	if _, err := exec.LookPath("pandoc"); err != nil {
 109  		return nil, fmt.Errorf("extract: pandoc not found: %w", err)
 110  	}
 111  	cmd := exec.Command("pandoc", "-t", "plain", "--wrap=none", path)
 112  	stdout, err := cmd.StdoutPipe()
 113  	if err != nil {
 114  		return nil, err
 115  	}
 116  	if err := cmd.Start(); err != nil {
 117  		return nil, err
 118  	}
 119  	return &cmdReadCloser{ReadCloser: stdout, cmd: cmd}, nil
 120  }
 121  
 122  // cmdReadCloser wraps a command's stdout pipe and waits for the command
 123  // on Close.
 124  type cmdReadCloser struct {
 125  	io.ReadCloser
 126  	cmd *exec.Cmd
 127  }
 128  
 129  func (c *cmdReadCloser) Close() error {
 130  	err := c.ReadCloser.Close()
 131  	if werr := c.cmd.Wait(); werr != nil && err == nil {
 132  		err = werr
 133  	}
 134  	return err
 135  }
 136  
 137  // IsTextFile does a quick heuristic check: reads the first 512 bytes
 138  // and checks if they look like text (no null bytes, mostly printable).
 139  func IsTextFile(path string) bool {
 140  	f, err := os.Open(path)
 141  	if err != nil {
 142  		return false
 143  	}
 144  	defer f.Close()
 145  
 146  	buf := make([]byte, 512)
 147  	n, err := f.Read(buf)
 148  	if err != nil && err != io.EOF {
 149  		return false
 150  	}
 151  	buf = buf[:n]
 152  	for _, b := range buf {
 153  		if b == 0 {
 154  			return false
 155  		}
 156  	}
 157  	return true
 158  }
 159