feat(downloader): allow download of grouped message. (iyear#795)

* feat(downloader): allow download of grouped message. * feat(dl): support auto detect grouped message * docs(dl): support auto detect grouped message * lint * fix iter elem id * fix unique iter elem id --------- Co-authored-by: iyear <[email protected]>
thejimmylin · Nov 17, 2024 · bdd84cc · bdd84cc
1 parent 2959e0d
commit bdd84cc
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 7 deletions.
diff --git a/app/dl/dl.go b/app/dl/dl.go
@@ -38,6 +38,7 @@ type Options struct {
 	Exclude    []string
 	Desc       bool
 	Takeout    bool
+	Group      bool // auto detect grouped message
 
 	// resume opts
 	Continue, Restart bool

diff --git a/app/dl/iter.go b/app/dl/iter.go
@@ -15,6 +15,8 @@ import (
 
 	"github.com/go-faster/errors"
 	"github.com/gotd/td/telegram/peers"
+	"github.com/gotd/td/tg"
+	"go.uber.org/atomic"
 
 	"github.com/iyear/tdl/core/dcpool"
 	"github.com/iyear/tdl/core/downloader"
@@ -53,7 +55,8 @@ type iter struct {
 	fingerprint string
 	preSum      []int
 	i, j        int
-	elem        downloader.Elem
+	counter     *atomic.Int64
+	elem        chan downloader.Elem
 	err         error
 }
 
@@ -96,7 +99,8 @@ func newIter(pool dcpool.Pool, manager *peers.Manager, dialog [][]*tmessage.Dial
 		preSum:      preSum(dialogs),
 		i:           0,
 		j:           0,
-		elem:        nil,
+		counter:     atomic.NewInt64(-1),
+		elem:        make(chan downloader.Elem, 10), // grouped message buffer
 		err:         nil,
 	}, nil
 }
@@ -113,6 +117,9 @@ func (i *iter) Next(ctx context.Context) bool {
 	if i.delay > 0 && (i.i+i.j) > 0 { // skip first delay
 		time.Sleep(i.delay)
 	}
+	if len(i.elem) > 0 { // there are messages(grouped) in channel that not processed
+		return true
+	}
 
 	for {
 		ok, skip := i.process(ctx)
@@ -152,13 +159,19 @@ func (i *iter) process(ctx context.Context) (ret bool, skip bool) {
 		i.err = errors.Wrap(err, "resolve from input peer")
 		return false, false
 	}
-
 	message, err := tutil.GetSingleMessage(ctx, i.pool.Default(ctx), peer, msg)
 	if err != nil {
 		i.err = errors.Wrap(err, "resolve message")
 		return false, false
 	}
 
+	if _, ok := message.GetGroupedID(); ok && i.opts.Group {
+		return i.processGrouped(ctx, message, from)
+	}
+	return i.processSingle(message, from)
+}
+
+func (i *iter) processSingle(message *tg.Message, from peers.Peer) (bool, bool) {
 	item, ok := tmedia.GetMedia(message)
 	if !ok {
 		i.err = errors.Errorf("can not get media from %d/%d message", from.ID(), message.ID)
@@ -175,7 +188,7 @@ func (i *iter) process(ctx context.Context) (ret bool, skip bool) {
 	}
 
 	toName := bytes.Buffer{}
-	err = i.tpl.Execute(&toName, &fileTemplate{
+	err := i.tpl.Execute(&toName, &fileTemplate{
 		DialogID:     from.ID(),
 		MessageID:    message.ID,
 		MessageDate:  int64(message.Date),
@@ -213,8 +226,8 @@ func (i *iter) process(ctx context.Context) (ret bool, skip bool) {
 		return false, false
 	}
 
-	i.elem = &iterElem{
-		id: i.ij2n(i.i, i.j),
+	i.elem <- &iterElem{
+		id: int(i.counter.Inc()),
 
 		from:    from,
 		fromMsg: message,
@@ -228,8 +241,22 @@ func (i *iter) process(ctx context.Context) (ret bool, skip bool) {
 	return true, false
 }
 
+func (i *iter) processGrouped(ctx context.Context, message *tg.Message, from peers.Peer) (bool, bool) {
+	grouped, err := tutil.GetGroupedMessages(ctx, i.pool.Default(ctx), from.InputPeer(), message)
+	if err != nil {
+		i.err = errors.Wrapf(err, "resolve grouped message %d/%d", from.ID(), message.ID)
+		return false, false
+	}
+
+	for _, msg := range grouped {
+		// best effort, ignore error
+		_, _ = i.processSingle(msg, from)
+	}
+	return true, false
+}
+
 func (i *iter) Value() downloader.Elem {
-	return i.elem
+	return <-i.elem
 }
 
 func (i *iter) Err() error {

diff --git a/cmd/dl.go b/cmd/dl.go
@@ -67,6 +67,7 @@ func NewDownload() *cobra.Command {
 
 	cmd.Flags().BoolVar(&opts.Desc, "desc", false, "download files from the newest to the oldest ones (may affect resume download)")
 	cmd.Flags().BoolVar(&opts.Takeout, "takeout", false, "takeout sessions let you export data from your account with lower flood wait limits.")
+	cmd.Flags().BoolVar(&opts.Group, "group", false, "auto detect grouped message and download all of them")
 
 	// resume flags, if both false then ask user
 	cmd.Flags().BoolVar(&opts.Continue, _continue, false, "continue the last download directly")

diff --git a/docs/content/en/guide/download.md b/docs/content/en/guide/download.md
@@ -90,6 +90,14 @@ Side effect: like `.apk` file, it will be renamed to `.zip`.
 tdl dl -u https://t.me/tdl/1 --rewrite-ext
 {{< /command >}}
 
+## Album/Grouped Detection
+
+Automatically detect if the message is an album/grouped message and download all of them.
+
+{{< command >}}
+tdl dl -u https://t.me/tdl/1 --group
+{{< /command >}}
+
 ## Auto Skip
 
 Skip the same files(name and size) when downloading.

diff --git a/docs/content/zh/guide/download.md b/docs/content/zh/guide/download.md
@@ -87,6 +87,14 @@ tdl dl -f result.json --desc
 tdl dl -u https://t.me/tdl/1 --rewrite-ext
 {{< /command >}}
 
+## 相册/组合消息探测
+
+自动检测消息是否为相册/组合消息，并下载相应的所有文件。
+
+{{< command >}}
+tdl dl -u https://t.me/tdl/1 --group
+{{< /command >}}
+
 ## 自动跳过
 
 在下载时跳过相同的文件（即名称和大小相同）。