mirror of
https://github.com/bit1001/tdl.git
synced 2024-09-20 14:45:55 +08:00
113 lines
2.3 KiB
Go
113 lines
2.3 KiB
Go
package downloader
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"github.com/fatih/color"
|
|
"github.com/gabriel-vasile/mimetype"
|
|
"github.com/gotd/td/telegram/downloader"
|
|
"github.com/gotd/td/tg"
|
|
"github.com/iyear/tdl/pkg/consts"
|
|
"github.com/iyear/tdl/pkg/prog"
|
|
"github.com/iyear/tdl/pkg/utils"
|
|
"github.com/jedib0t/go-pretty/v6/progress"
|
|
"golang.org/x/sync/errgroup"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const TempExt = ".tmp"
|
|
|
|
type Downloader struct {
|
|
client *tg.Client
|
|
pw progress.Writer
|
|
partSize int
|
|
threads int
|
|
iter Iter
|
|
}
|
|
|
|
func New(client *tg.Client, partSize int, threads int, iter Iter) *Downloader {
|
|
return &Downloader{
|
|
client: client,
|
|
pw: prog.New(),
|
|
partSize: partSize,
|
|
threads: threads,
|
|
iter: iter,
|
|
}
|
|
}
|
|
|
|
func (d *Downloader) Download(ctx context.Context, limit int) error {
|
|
d.pw.SetNumTrackersExpected(d.iter.Total(ctx))
|
|
|
|
go d.pw.Render()
|
|
|
|
wg, errctx := errgroup.WithContext(ctx)
|
|
wg.SetLimit(limit)
|
|
|
|
for d.iter.Next(ctx) {
|
|
item, err := d.iter.Value(ctx)
|
|
if err != nil {
|
|
d.pw.Log(color.RedString("[ERROR] Get item failed: %v", err))
|
|
continue
|
|
}
|
|
|
|
wg.Go(func() error {
|
|
// d.pw.Log(color.MagentaString("name: %s,size: %s", item.Name, utils.Byte.FormatBytes(item.Size)))
|
|
return d.download(errctx, item)
|
|
})
|
|
}
|
|
|
|
err := wg.Wait()
|
|
if err != nil {
|
|
|
|
return err
|
|
}
|
|
|
|
for d.pw.IsRenderInProgress() {
|
|
if d.pw.LengthActive() == 0 {
|
|
d.pw.Stop()
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (d *Downloader) download(ctx context.Context, item *Item) error {
|
|
tracker := prog.AppendTracker(d.pw, item.Name, item.Size)
|
|
filename := fmt.Sprintf("%s%s", utils.FS.GetNameWithoutExt(item.Name), TempExt)
|
|
path := filepath.Join(consts.DownloadPath, filename)
|
|
|
|
f, err := os.Create(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = downloader.NewDownloader().WithPartSize(d.partSize).
|
|
Download(d.client, item.InputFileLoc).WithThreads(d.threads).
|
|
Parallel(ctx, &writeAt{
|
|
f: f,
|
|
tracker: tracker,
|
|
})
|
|
if err = f.Close(); err != nil {
|
|
return err
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
mime, err := mimetype.DetectFile(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
newfile := fmt.Sprintf("%s%s", strings.TrimSuffix(filename, TempExt), mime.Extension())
|
|
if err = os.Rename(path, filepath.Join(consts.DownloadPath, newfile)); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|