yt-download.go

package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/knadh/go-get-youtube/youtube"

	"github.com/pkg/errors"
)

var (
	csvfile = flag.String("csv", "", "a csv file to pull titles and youtube urls from")
	outdir  = flag.String("outdir", ".", "the directory all files should be output to")
)

func main() {
	flag.Parse()
	if csvfile == nil || *csvfile == "" {
		// Todo: support alternatives
		fmt.Println("Please provide a value for --csv")
		return
	}
	f, err := os.Open(*csvfile)
	defer f.Close()
	if err != nil {
		fmt.Println("Failed to open csvfile")
	}
	r := csv.NewReader(f)
	r.ReuseRecord = true
	r.FieldsPerRecord = 2

	// Prepare already-downloaded files
	dir, err := os.Open(*outdir)
	if err != nil {
		fmt.Println(err)
		return
	}
	names, err := dir.Readdirnames(0)
	if err != nil {
		fmt.Println(err)
		return
	}
	nameMap := make(map[string]struct{}, len(names))
	for _, n := range names {
		nameMap[n] = struct{}{}
	}

	wg := sync.WaitGroup{}
	for {
		fmt.Println("Reading record")
		records, err := r.Read()
		if err == io.EOF {
			break
		} else if err == csv.ErrFieldCount {
			fmt.Println("Expected the csv to be a list of 'song title, url' pairs")
			return
		} else if err != nil {
			fmt.Println("Unexpected error reading csv", err)
			return
		}
		req := DownloadReq{
			Outfile: records[0],
			URL:     records[1],
			Audio:   true,
			Outdir:  *outdir,
		}
		wg.Add(1)
		go func(req DownloadReq) {
			defer wg.Done()
			// If we already have this file, skip
			if _, ok := nameMap[req.Outfile+".mp3"]; ok {
				fmt.Println("Already have", req.Outfile)
				return
			}
			err := req.Download()
			if err != nil {
				fmt.Println("Error downloading file", req, err)
			}
		}(req)
	}
	wg.Wait()
}

type DownloadReq struct {
	Outfile string
	Outdir  string
	URL     string
	// Just output audio, not video
	Audio bool
}

func (dr DownloadReq) Download() error {
	fullpath := filepath.Join(dr.Outdir, dr.Outfile)
	// We anticipate a string of the format "...youtube.com/watch?v=CODE"
	match := "youtube.com/watch?v="
	idx := strings.Index(dr.URL, match)
	if idx == -1 {
		return errors.New("malformed url has no " + match + " :" + dr.URL)
	}
	// Get to the code
	idx += len(match)
	code := dr.URL[idx:]
	// Request file
	metadata, err := youtube.Get(code)
	if err != nil {
		return errors.Wrap(err, "failed to query youtube for metadata")
	}
	opts := &youtube.Option{
		Mp3: dr.Audio,
	}
	fmt.Println(len(metadata.Formats), metadata.Formats[0])
	metadata.Download(0, fullpath, opts)
	return nil
}