From 5233607250bd95fbfc66aa2bf8166e4489d634f7 Mon Sep 17 00:00:00 2001 From: iyear Date: Mon, 27 Nov 2023 11:03:41 +0800 Subject: [PATCH] refactor(up): extract to interface --- app/up/elem.go | 52 +++++++++++++++++++++++++++ app/up/iter.go | 76 +++++++++++----------------------------- app/up/progress.go | 60 ++++++++++++++++++++++++------- app/up/up.go | 1 + pkg/uploader/iter.go | 23 ++++++------ pkg/uploader/progress.go | 8 ++--- pkg/uploader/uploader.go | 63 ++++++++++++++++++--------------- 7 files changed, 169 insertions(+), 114 deletions(-) create mode 100644 app/up/elem.go diff --git a/app/up/elem.go b/app/up/elem.go new file mode 100644 index 0000000..8c2d0a0 --- /dev/null +++ b/app/up/elem.go @@ -0,0 +1,52 @@ +package up + +import ( + "os" + "path/filepath" + + "github.com/gotd/td/telegram/peers" + "github.com/gotd/td/tg" + + "github.com/iyear/tdl/pkg/uploader" +) + +type iterElem struct { + file *uploaderFile + thumb *uploaderFile + to peers.Peer + + asPhoto bool + remove bool +} + +func (e *iterElem) File() uploader.File { + return e.file +} + +func (e *iterElem) Thumb() (uploader.File, bool) { + if e.thumb == nil { + return nil, false + } + return e.thumb, true +} + +func (e *iterElem) To() tg.InputPeerClass { + return e.to.InputPeer() +} + +func (e *iterElem) AsPhoto() bool { + return e.asPhoto +} + +type uploaderFile struct { + *os.File + size int64 +} + +func (u *uploaderFile) Name() string { + return filepath.Base(u.File.Name()) +} + +func (u *uploaderFile) Size() int64 { + return u.size +} diff --git a/app/up/iter.go b/app/up/iter.go index 0e9b6e4..da2f251 100644 --- a/app/up/iter.go +++ b/app/up/iter.go @@ -3,7 +3,6 @@ package up import ( "context" "os" - "path/filepath" "github.com/gabriel-vasile/mimetype" "github.com/go-faster/errors" @@ -26,18 +25,19 @@ type iter struct { cur int err error - file *uploader.Elem + file uploader.Elem } func newIter(files []*file, to peers.Peer, photo, remove bool) *iter { return &iter{ files: files, - cur: 0, - err: nil, - file: nil, to: to, photo: photo, remove: remove, + + cur: 0, + err: nil, + file: nil, } } @@ -56,25 +56,13 @@ func (i *iter) Next(ctx context.Context) bool { cur := i.files[i.cur] i.cur++ - fMime, err := mimetype.DetectFile(cur.file) - if err != nil { - i.err = errors.Wrap(err, "detect mime") - return false - } - f, err := os.Open(cur.file) if err != nil { i.err = errors.Wrap(err, "open file") return false } - stat, err := f.Stat() - if err != nil { - i.err = errors.Wrap(err, "stat file") - return false - } - - var thumb uploader.File = nopFile{} + var thumb *uploaderFile = nil // has thumbnail if cur.thumb != "" { tMime, err := mimetype.DetectFile(cur.thumb) @@ -88,53 +76,31 @@ func (i *iter) Next(ctx context.Context) bool { return false } - thumb = uploaderFile{thumbFile} + thumb = &uploaderFile{File: thumbFile, size: 0} } - i.file = &uploader.Elem{ - File: uploaderFile{f}, - Thumb: thumb, - Name: filepath.Base(f.Name()), - MIME: fMime.String(), - Size: stat.Size(), - To: i.to, - Photo: i.photo, - Remove: i.remove, + stat, err := f.Stat() + if err != nil { + i.err = errors.Wrap(err, "stat file") + return false + } + + i.file = &iterElem{ + file: &uploaderFile{File: f, size: stat.Size()}, + thumb: thumb, + to: i.to, + + asPhoto: i.photo, + remove: i.remove, } return true } -func (i *iter) Value() *uploader.Elem { +func (i *iter) Value() uploader.Elem { return i.file } func (i *iter) Err() error { return i.err } - -type nopFile struct{} - -func (nopFile) Read(_ []byte) (n int, err error) { - return 0, errors.New("nopFile") -} - -func (nopFile) Seek(_ int64, _ int) (int64, error) { - return 0, errors.New("nopFile") -} - -func (nopFile) Close() error { - return nil -} - -func (nopFile) Remove() error { - return nil -} - -type uploaderFile struct { - *os.File -} - -func (u uploaderFile) Remove() error { - return os.Remove(u.Name()) -} diff --git a/app/up/progress.go b/app/up/progress.go index 88ac1b0..0145131 100644 --- a/app/up/progress.go +++ b/app/up/progress.go @@ -2,9 +2,11 @@ package up import ( "fmt" + "os" "sync" "github.com/fatih/color" + "github.com/go-faster/errors" pw "github.com/jedib0t/go-pretty/v6/progress" "github.com/iyear/tdl/pkg/prog" @@ -29,12 +31,12 @@ func newProgress(p pw.Writer) *progress { } } -func (p *progress) OnAdd(elem *uploader.Elem) { - tracker := prog.AppendTracker(p.pw, utils.Byte.FormatBinaryBytes, p.processMessage(elem), elem.Size) +func (p *progress) OnAdd(elem uploader.Elem) { + tracker := prog.AppendTracker(p.pw, utils.Byte.FormatBinaryBytes, p.processMessage(elem), elem.File().Size()) p.trackers.Store(p.tuple(elem), tracker) } -func (p *progress) OnUpload(elem *uploader.Elem, state uploader.ProgressState) { +func (p *progress) OnUpload(elem uploader.Elem, state uploader.ProgressState) { tracker, ok := p.trackers.Load(p.tuple(elem)) if !ok { return @@ -45,28 +47,60 @@ func (p *progress) OnUpload(elem *uploader.Elem, state uploader.ProgressState) { t.SetValue(state.Uploaded) } -func (p *progress) OnDone(elem *uploader.Elem, err error) { +func (p *progress) OnDone(elem uploader.Elem, err error) { tracker, ok := p.trackers.Load(p.tuple(elem)) if !ok { return } - t := tracker.(*pw.Tracker) - if err != nil { - p.pw.Log(color.RedString("%s error: %s", p.elemString(elem), err.Error())) - t.MarkAsErrored() + e := elem.(*iterElem) + + if err := p.closeFile(e); err != nil { + p.fail(t, elem, errors.Wrap(err, "close file")) return } + + if err != nil { + p.fail(t, elem, errors.Wrap(err, "progress")) + return + } + + if e.remove { + if err := os.Remove(e.file.File.Name()); err != nil { + p.fail(t, elem, errors.Wrap(err, "remove file")) + return + } + } } -func (p *progress) tuple(elem *uploader.Elem) tuple { - return tuple{elem.Name, elem.To.ID()} +func (p *progress) closeFile(e *iterElem) error { + if err := e.file.Close(); err != nil { + return errors.Wrap(err, "close file") + } + + if e.thumb != nil { + if err := e.thumb.Close(); err != nil { + return errors.Wrap(err, "close thumb") + } + } + + return nil } -func (p *progress) processMessage(elem *uploader.Elem) string { +func (p *progress) fail(t *pw.Tracker, elem uploader.Elem, err error) { + p.pw.Log(color.RedString("%s error: %s", p.elemString(elem), err.Error())) + t.MarkAsErrored() +} + +func (p *progress) tuple(elem uploader.Elem) tuple { + return tuple{elem.(*iterElem).file.File.Name(), elem.(*iterElem).to.ID()} +} + +func (p *progress) processMessage(elem uploader.Elem) string { return p.elemString(elem) } -func (p *progress) elemString(elem *uploader.Elem) string { - return fmt.Sprintf("%s -> %s(%d)", elem.Name, elem.To.VisibleName(), elem.To.ID()) +func (p *progress) elemString(elem uploader.Elem) string { + e := elem.(*iterElem) + return fmt.Sprintf("%s -> %s(%d)", e.file.File.Name(), e.to.VisibleName(), e.to.ID()) } diff --git a/app/up/up.go b/app/up/up.go index 5c8131b..9ebb098 100644 --- a/app/up/up.go +++ b/app/up/up.go @@ -57,6 +57,7 @@ func Run(ctx context.Context, opts *Options) error { } upProgress := prog.New(utils.Byte.FormatBinaryBytes) + upProgress.SetNumTrackersExpected(len(files)) prog.EnablePS(ctx, upProgress) options := uploader.Options{ diff --git a/pkg/uploader/iter.go b/pkg/uploader/iter.go index bb9464f..8e0d778 100644 --- a/pkg/uploader/iter.go +++ b/pkg/uploader/iter.go @@ -4,27 +4,24 @@ import ( "context" "io" - "github.com/gotd/td/telegram/peers" + "github.com/gotd/td/tg" ) type Iter interface { Next(ctx context.Context) bool - Value() *Elem + Value() Elem Err() error } type File interface { - io.ReadSeekCloser - Remove() error + io.ReadSeeker + Name() string + Size() int64 } -type Elem struct { - File File - Thumb File - Name string - MIME string - Size int64 - To peers.Peer - Photo bool - Remove bool +type Elem interface { + File() File + Thumb() (File, bool) + To() tg.InputPeerClass + AsPhoto() bool } diff --git a/pkg/uploader/progress.go b/pkg/uploader/progress.go index 866f93d..22e23a4 100644 --- a/pkg/uploader/progress.go +++ b/pkg/uploader/progress.go @@ -7,9 +7,9 @@ import ( ) type Progress interface { - OnAdd(elem *Elem) - OnUpload(elem *Elem, state ProgressState) - OnDone(elem *Elem, err error) + OnAdd(elem Elem) + OnUpload(elem Elem, state ProgressState) + OnDone(elem Elem, err error) // TODO: OnLog to log something that is not an error but should be sent to the user } @@ -19,7 +19,7 @@ type ProgressState struct { } type wrapProcess struct { - elem *Elem + elem Elem process Progress } diff --git a/pkg/uploader/uploader.go b/pkg/uploader/uploader.go index 841e4be..95a7d40 100644 --- a/pkg/uploader/uploader.go +++ b/pkg/uploader/uploader.go @@ -2,16 +2,15 @@ package uploader import ( "context" - "fmt" "io" "time" + "github.com/gabriel-vasile/mimetype" "github.com/go-faster/errors" "github.com/gotd/td/telegram/message" "github.com/gotd/td/telegram/message/styling" "github.com/gotd/td/telegram/uploader" "github.com/gotd/td/tg" - "go.uber.org/multierr" "golang.org/x/sync/errgroup" "github.com/iyear/tdl/pkg/utils" @@ -64,17 +63,7 @@ func (u *Uploader) Upload(ctx context.Context, limit int) error { return wg.Wait() } -func (u *Uploader) upload(ctx context.Context, elem *Elem) (rerr error) { - defer func() { - if rerr == nil && elem.Remove { - multierr.AppendInto(&rerr, elem.File.Remove()) - multierr.AppendInto(&rerr, elem.Thumb.Remove()) - } - }() - - defer multierr.AppendInvoke(&rerr, multierr.Close(elem.File)) - defer multierr.AppendInvoke(&rerr, multierr.Close(elem.Thumb)) - +func (u *Uploader) upload(ctx context.Context, elem Elem) error { select { case <-ctx.Done(): return ctx.Err() @@ -89,48 +78,64 @@ func (u *Uploader) upload(ctx context.Context, elem *Elem) (rerr error) { process: u.opts.Progress, }) - f, err := up.Upload(ctx, uploader.NewUpload(elem.Name, elem.File, elem.Size)) + f, err := up.Upload(ctx, uploader.NewUpload(elem.File().Name(), elem.File(), elem.File().Size())) if err != nil { return errors.Wrap(err, "upload file") } - caption := []message.StyledTextOption{ - styling.Code(elem.Name), - styling.Plain(" - "), - styling.Code(elem.MIME), + if _, err = elem.File().Seek(0, io.SeekStart); err != nil { + return errors.Wrap(err, "seek file") } - doc := message.UploadedDocument(f, caption...).MIME(elem.MIME).Filename(elem.Name) + mime, err := mimetype.DetectReader(elem.File()) + if err != nil { + return errors.Wrap(err, "detect mime") + } + + caption := []message.StyledTextOption{ + styling.Code(elem.File().Name()), + styling.Plain(" - "), + styling.Code(mime.String()), + } + doc := message.UploadedDocument(f, caption...). + MIME(mime.String()). + Filename(elem.File().Name()) // upload thumbnail TODO(iyear): maybe still unavailable - if thumb, err := uploader.NewUploader(u.opts.Client). - FromReader(ctx, fmt.Sprintf("%s.thumb", elem.Name), elem.Thumb); err == nil { - doc = doc.Thumb(thumb) + if thumb, ok := elem.Thumb(); ok { + if thumbFile, err := uploader.NewUploader(u.opts.Client). + FromReader(ctx, thumb.Name(), thumb); err == nil { + doc = doc.Thumb(thumbFile) + } } var media message.MediaOption = doc switch { - case utils.Media.IsImage(elem.MIME) && elem.Photo: + case utils.Media.IsImage(mime.String()) && elem.AsPhoto(): + // webp should be uploaded as document + if mime.String() == "image/webp" { + break + } // upload as photo media = message.UploadedPhoto(f, caption...) - case utils.Media.IsVideo(elem.MIME): + case utils.Media.IsVideo(mime.String()): // reset reader - if _, err = elem.File.Seek(0, io.SeekStart); err != nil { + if _, err = elem.File().Seek(0, io.SeekStart); err != nil { return errors.Wrap(err, "seek file") } - if dur, w, h, err := utils.Media.GetMP4Info(elem.File); err == nil { + if dur, w, h, err := utils.Media.GetMP4Info(elem.File()); err == nil { // #132. There may be some errors, but we can still upload the file media = doc.Video(). Duration(time.Duration(dur)*time.Second). Resolution(w, h). SupportsStreaming() } - case utils.Media.IsAudio(elem.MIME): - media = doc.Audio().Title(utils.FS.GetNameWithoutExt(elem.Name)) + case utils.Media.IsAudio(mime.String()): + media = doc.Audio().Title(utils.FS.GetNameWithoutExt(elem.File().Name())) } _, err = message.NewSender(u.opts.Client). WithUploader(up). - To(elem.To.InputPeer()). + To(elem.To()). Media(ctx, media) if err != nil { return errors.Wrap(err, "send message")