package ytdlp import ( "context" "encoding/json" "io" "os/exec" "strings" "github.com/royalcat/ctxprogress" "golang.org/x/sync/errgroup" ) type DownloadStatus string const ( StatusDownloading DownloadStatus = "downloading" StatusFinished DownloadStatus = "finished" StatusErrored DownloadStatus = "error" ) // Progress for the Running call type DownloadProgress struct { Status DownloadStatus `json:"status"` Filename string `json:"filename"` TmpFilename string `json:"tmpfilename"` DownloadedBytes int64 `json:"downloaded_bytes"` TotalBytes int64 `json:"total_bytes"` TotalBytesEstimate float64 `json:"total_bytes_estimate"` Elapsed float64 `json:"elapsed"` ETA float64 `json:"eta"` Speed float64 `json:"speed"` FragmentIndex int64 `json:"fragment_index"` FragmentCount int64 `json:"fragment_count"` } // Current implements ctxprogress.Progress. func (d DownloadProgress) Progress() (int, int) { if d.TotalBytes != -1 && d.TotalBytes != 0 && d.DownloadedBytes != -1 { return int(d.DownloadedBytes), int(d.TotalBytes) } if d.TotalBytesEstimate != -1 && d.TotalBytesEstimate != 0 && d.DownloadedBytes != -1 { return int(d.DownloadedBytes), int(d.TotalBytesEstimate) } return int(d.FragmentIndex), int(d.FragmentCount) } const rawProgressTemplate = `download: %{ "status":"%(progress.status)s", "eta":%(progress.eta|-1)s, "speed":%(progress.speed|0)s, "downloaded_bytes":%(progress.downloaded_bytes|-1)s, "total_bytes": %(progress.total_bytes|-1)s, "total_bytes_estimate": %(progress.total_bytes_estimate|-1)s, "fragment_index":%(progress.fragment_index|-1)s, "fragment_count":%(progress.fragment_count|-1)s }` var progressTemplate = strings.NewReplacer("\n", "", "\t", "", " ", "").Replace(rawProgressTemplate) func (c *Client) Download(ctx context.Context, url string, w io.Writer) error { args := []string{ "--progress", "--newline", "--progress-template", progressTemplate, "-o", "-", url, } group, ctx := errgroup.WithContext(ctx) stderr, lines, err := lineReader(group) if err != nil { return err } cmd := exec.CommandContext(ctx, c.binary, args...) cmd.Stdout = w cmd.Stderr = stderr group.Go(func() error { err := cmd.Run() stderr.Close() if err != nil { return err } return nil }) for line := range lines { if line, ok := strings.CutPrefix(line, "%"); ok { p := DownloadProgress{} err = json.Unmarshal([]byte(line), &p) if err != nil { //TODO: handle error continue } ctxprogress.Set(ctx, p) } } return group.Wait() }