packer-cn/cmd/snippet-extractor/main.go

// snippet-extract --begin=#_BEGIN_WRAP_TAG_ --end=#_END_WRAP_TAG_ -output_dir=./docs/ [./file|./directory/]...
// Extracts markdown snippets from relative files into output_dir, keeping the
// directory layout.
// It is not mandatory to terminate a snippet, the extractor will simply add
// line until EOF.
// Lines matching begin or end tags will not be put in the resulting file.
// When a directory is passed, all files from directory will be parsed.
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

type options struct {
	filenames  []string
	begin, end string
	outputDir  string
	extension  string
}

func (o *options) AddFlagSets(fs *flag.FlagSet) {
	fs.StringVar(&o.begin, "begin", "#_BEGIN_WRAP_TAG_", "flag to mark beginning of a snippet")
	fs.StringVar(&o.end, "end", "#_END_WRAP_TAG_", "flag to mark ending of a snippet")
	fs.StringVar(&o.outputDir, "output_dir", "./docs/", "flag to mark ending of a snippet")
	fs.StringVar(&o.extension, "extension", ".mdx", "extension for generated files")
}

func main() {
	fs := flag.NewFlagSet("snippet-extract", flag.ContinueOnError)
	opts := options{}
	opts.AddFlagSets(fs)
	if err := fs.Parse(os.Args[1:]); err != nil {
		fs.Usage()
		os.Exit(1)
	}
	wd, _ := os.Getwd()
	fmt.Printf("working from %s\n", wd)
	opts.filenames = extactFilenames(fs.Args())
	for _, filename := range opts.filenames {
		ext := filepath.Ext(filename)
		snippets := extractSnippets(opts.begin, opts.end, filename)
		for _, snippet := range snippets {
			outputFile := filepath.Join(opts.outputDir, filepath.Base(filename), snippet.Identifier+opts.extension)
			folder := filepath.Dir(outputFile)
			err := os.MkdirAll(folder, os.ModePerm)
			if err != nil {
				log.Printf("cannot mkdir %s: %s", folder, err)
			}
			f := bytes.NewBuffer(nil)
			fmt.Fprintf(f, `<!-- Code generated by snippet-extractor %s; DO NOT EDIT MANUALLY -->`, strings.Join(os.Args[1:], " "))
			fmt.Fprintf(f, "\n\n```%s\n%s```\n", ext, snippet.Text)
			err = ioutil.WriteFile(outputFile, f.Bytes(), 0600)
			if err != nil {
				log.Printf("cannot write %s in %s: %s", filepath.Base(outputFile), folder, err)
			}
		}
	}
}

type snippet struct {
	Identifier string
	Text       string
	Closed     bool
}

func extractSnippets(beginPattern, endPattern, filename string) []snippet {
	file, err := os.Open(filename)
	if err != nil {
		log.Printf("could not open file: %s", err)
		os.Exit(1)
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	snippets := []snippet{}
	for scanner.Scan() {
		line := scanner.Text()
		if identifier := matches(line, beginPattern); identifier != "" {
			snippets = append(snippets, snippet{
				Identifier: identifier,
			})
			continue
		}
		if identifier := matches(line, endPattern); identifier != "" {
			for i := range snippets {
				snippet := &snippets[i]
				if snippet.Identifier == identifier {
					snippet.Closed = true
				}
			}
			continue
		}
		for i := range snippets {
			snippet := &snippets[i]
			if snippet.Closed {
				continue
			}
			snippet.Text = snippet.Text + line + "\n"
		}
	}
	return snippets
}

func matches(s, prefix string) string {
	trimmed := strings.TrimSpace(s)
	lenDiff := len(s) - len(trimmed)
	if strings.HasPrefix(trimmed, prefix) {
		return s[len(prefix)+lenDiff:]
	}
	return ""
}

// if an entry is a directory all files from directory will be listed.
func extactFilenames(in []string) []string {
	out := []string{}
	for _, path := range in {
		fi, err := os.Stat(path)
		if err != nil {
			log.Fatalf("%s: %s", path, err)
		}
		if !fi.IsDir() {
			out = append(out, path)
			continue
		}
		files, err := ioutil.ReadDir(path)
		if err != nil {
			log.Fatalf("could not read directory %s: %s", path, err)
		}
		for _, file := range files {
			if file.IsDir() {
				continue
			}
			out = append(out, file.Name())
		}
	}
	return in
}