tstor/pkg/ytdlp/download.go

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()
}