From e1de0f5bdbf5034b6c92e211a2b8c20e1e1dd975 Mon Sep 17 00:00:00 2001
From: XMLHexagram <xml.hexagram@gmail.com>
Date: Fri, 29 Mar 2024 14:06:29 +0800
Subject: [PATCH] feat(upload): support `--include` flag

---
 app/dl/iter.go             | 13 +++----------
 app/up/up.go               |  3 ++-
 app/up/walk.go             | 21 +++++++++++++--------
 cmd/up.go                  |  1 +
 pkg/filterMap/filterMap.go |  9 +++++++++
 5 files changed, 28 insertions(+), 19 deletions(-)
 create mode 100644 pkg/filterMap/filterMap.go

diff --git a/app/dl/iter.go b/app/dl/iter.go
index db96029a26..9546f3e42b 100644
--- a/app/dl/iter.go
+++ b/app/dl/iter.go
@@ -6,6 +6,7 @@ import (
 	"crypto/sha256"
 	"encoding/binary"
 	"fmt"
+	"github.com/iyear/tdl/pkg/filterMap"
 	"os"
 	"path/filepath"
 	"sort"
@@ -68,8 +69,8 @@ func newIter(pool dcpool.Pool, manager *peers.Manager, dialog [][]*tmessage.Dial
 	}
 
 	// include and exclude
-	includeMap := filterMap(opts.Include, utils.FS.AddPrefixDot)
-	excludeMap := filterMap(opts.Exclude, utils.FS.AddPrefixDot)
+	includeMap := filterMap.New(opts.Include, utils.FS.AddPrefixDot)
+	excludeMap := filterMap.New(opts.Exclude, utils.FS.AddPrefixDot)
 
 	// to keep fingerprint stable
 	sortDialogs(dialogs, opts.Desc)
@@ -274,14 +275,6 @@ func flatDialogs(dialogs [][]*tmessage.Dialog) []*tmessage.Dialog {
 	return res
 }
 
-func filterMap(data []string, keyFn func(key string) string) map[string]struct{} {
-	m := make(map[string]struct{})
-	for _, v := range data {
-		m[keyFn(v)] = struct{}{}
-	}
-	return m
-}
-
 func sortDialogs(dialogs []*tmessage.Dialog, desc bool) {
 	sort.Slice(dialogs, func(i, j int) bool {
 		return utils.Telegram.GetInputPeerID(dialogs[i].Peer) <
diff --git a/app/up/up.go b/app/up/up.go
index b63853917b..d25197dab9 100644
--- a/app/up/up.go
+++ b/app/up/up.go
@@ -24,13 +24,14 @@ import (
 type Options struct {
 	Chat     string
 	Paths    []string
+	Includes []string
 	Excludes []string
 	Remove   bool
 	Photo    bool
 }
 
 func Run(ctx context.Context, c *telegram.Client, kvd kv.KV, opts Options) (rerr error) {
-	files, err := walk(opts.Paths, opts.Excludes)
+	files, err := walk(opts.Paths, opts.Includes, opts.Excludes)
 	if err != nil {
 		return err
 	}
diff --git a/app/up/walk.go b/app/up/walk.go
index e493b2dd21..7b437b3d00 100644
--- a/app/up/walk.go
+++ b/app/up/walk.go
@@ -1,6 +1,7 @@
 package up
 
 import (
+	"github.com/iyear/tdl/pkg/filterMap"
 	"io/fs"
 	"path/filepath"
 	"strings"
@@ -9,15 +10,13 @@ import (
 	"github.com/iyear/tdl/pkg/utils"
 )
 
-func walk(paths, excludes []string) ([]*file, error) {
+func walk(paths, includes, excludes []string) ([]*file, error) {
 	files := make([]*file, 0)
-	excludesMap := map[string]struct{}{
-		consts.UploadThumbExt: {}, // ignore thumbnail files
-	}
 
-	for _, exclude := range excludes {
-		excludesMap[exclude] = struct{}{}
-	}
+	// include and exclude
+	includesMap := filterMap.New(includes, utils.FS.AddPrefixDot)
+	excludesMap := filterMap.New(excludes, utils.FS.AddPrefixDot)
+	excludesMap[consts.UploadThumbExt] = struct{}{} // ignore thumbnail files
 
 	for _, path := range paths {
 		err := filepath.WalkDir(path, func(path string, d fs.DirEntry, err error) error {
@@ -27,7 +26,13 @@ func walk(paths, excludes []string) ([]*file, error) {
 			if d.IsDir() {
 				return nil
 			}
-			if _, ok := excludesMap[filepath.Ext(path)]; ok {
+
+			// process include and exclude
+			ext := filepath.Ext(path)
+			if _, ok := includesMap[ext]; len(includesMap) > 0 && !ok {
+				return nil
+			}
+			if _, ok := excludesMap[ext]; len(excludesMap) > 0 && ok {
 				return nil
 			}
 
diff --git a/cmd/up.go b/cmd/up.go
index 17741fe0c6..0e57ed9b4a 100644
--- a/cmd/up.go
+++ b/cmd/up.go
@@ -31,6 +31,7 @@ func NewUpload() *cobra.Command {
 	)
 	cmd.Flags().StringVarP(&opts.Chat, _chat, "c", "", "chat id or domain, and empty means 'Saved Messages'")
 	cmd.Flags().StringSliceVarP(&opts.Paths, path, "p", []string{}, "dirs or files")
+	cmd.Flags().StringSliceVarP(&opts.Includes, "includes", "i", []string{}, "include the specified file extensions")
 	cmd.Flags().StringSliceVarP(&opts.Excludes, "excludes", "e", []string{}, "exclude the specified file extensions")
 	cmd.Flags().BoolVar(&opts.Remove, "rm", false, "remove the uploaded files after uploading")
 	cmd.Flags().BoolVar(&opts.Photo, "photo", false, "upload the image as a photo instead of a file")
diff --git a/pkg/filterMap/filterMap.go b/pkg/filterMap/filterMap.go
new file mode 100644
index 0000000000..5400c9e6ff
--- /dev/null
+++ b/pkg/filterMap/filterMap.go
@@ -0,0 +1,9 @@
+package filterMap
+
+func New(data []string, keyFn func(key string) string) map[string]struct{} {
+	m := make(map[string]struct{})
+	for _, v := range data {
+		m[keyFn(v)] = struct{}{}
+	}
+	return m
+}