Rewrite format filters and allow filtering by langauge

corny · corny · commit 2e06f1f71ee4 · 2023-12-24T22:11:31.000+01:00
diff --git a/client_test.go b/client_test.go
@@ -156,6 +156,30 @@ func TestGetVideoWithManifestURL(t *testing.T) {
 	assert.NotZero(size)
 }
 
+func TestGetVideo_MultiLanguage(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+	video, err := testClient.GetVideo("https://www.youtube.com/watch?v=pU9sHwNKc2c")
+	require.NoError(err)
+	require.NotNil(video)
+
+	// collect languages
+	var languageNames, lanaguageIDs []string
+	for _, format := range video.Formats {
+		if format.AudioTrack != nil {
+			languageNames = append(languageNames, format.LanguageDisplayName())
+			lanaguageIDs = append(lanaguageIDs, format.AudioTrack.ID)
+		}
+	}
+
+	assert.Contains(languageNames, "English original")
+	assert.Contains(languageNames, "Portuguese (Brazil)")
+	assert.Contains(lanaguageIDs, "en.4")
+	assert.Contains(lanaguageIDs, "pt-BR.3")
+
+	assert.Empty(video.Formats.Language("Does not exist"))
+	assert.NotEmpty(video.Formats.Language("English original"))
+}
+
 func TestGetStream(t *testing.T) {
 	assert, require := assert.New(t), require.New(t)
 
diff --git a/cmd/youtubedr/download.go b/cmd/youtubedr/download.go
@@ -32,8 +32,7 @@ func init() {
 
 	downloadCmd.Flags().StringVarP(&outputFile, "filename", "o", "", "The output file, the default is genated by the video title.")
 	downloadCmd.Flags().StringVarP(&outputDir, "directory", "d", ".", "The output directory.")
-	addQualityFlag(downloadCmd.Flags())
-	addMimeTypeFlag(downloadCmd.Flags())
+	addVideoSelectionFlags(downloadCmd.Flags())
 }
 
 func download(id string) error {
@@ -48,7 +47,7 @@ func download(id string) error {
 		if err := checkFFMPEG(); err != nil {
 			return err
 		}
-		return downloader.DownloadComposite(context.Background(), outputFile, video, outputQuality, mimetype)
+		return downloader.DownloadComposite(context.Background(), outputFile, video, outputQuality, mimetype, language)
 	}
 
 	return downloader.Download(context.Background(), video, format, outputFile)
diff --git a/cmd/youtubedr/downloader.go b/cmd/youtubedr/downloader.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"crypto/tls"
-	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -20,16 +19,15 @@ import (
 var (
 	insecureSkipVerify bool   // skip TLS server validation
 	outputQuality      string // itag number or quality string
-	mimetype           string // mimetype
+	mimetype           string
+	language           string
 	downloader         *ytdl.Downloader
 )
 
-func addQualityFlag(flagSet *pflag.FlagSet) {
+func addVideoSelectionFlags(flagSet *pflag.FlagSet) {
 	flagSet.StringVarP(&outputQuality, "quality", "q", "medium", "The itag number or quality label (hd720, medium)")
-}
-
-func addMimeTypeFlag(flagSet *pflag.FlagSet) {
-	flagSet.StringVarP(&mimetype, "mimetype", "m", "mp4", "Mime-Type to filter (mp4, webm, av01, avc1) - applicable if --quality used is quality label")
+	flagSet.StringVarP(&mimetype, "mimetype", "m", "", "Mime-Type to filter (mp4, webm, av01, avc1) - applicable if --quality used is quality label")
+	flagSet.StringVarP(&language, "language", "l", "", "Language to filter")
 }
 
 func getDownloader() *ytdl.Downloader {
@@ -70,41 +68,34 @@ func getDownloader() *ytdl.Downloader {
 	return downloader
 }
 
-func getVideoWithFormat(id string) (*youtube.Video, *youtube.Format, error) {
+func getVideoWithFormat(videoID string) (*youtube.Video, *youtube.Format, error) {
 	dl := getDownloader()
-	video, err := dl.GetVideo(id)
+	video, err := dl.GetVideo(videoID)
 	if err != nil {
 		return nil, nil, err
 	}
+
+	itag, _ := strconv.Atoi(outputQuality)
 	formats := video.Formats
+
+	if language != "" {
+		formats = formats.Language(language)
+	}
 	if mimetype != "" {
 		formats = formats.Type(mimetype)
 	}
-	if len(formats) == 0 {
-		return nil, nil, errors.New("no formats found")
+	if outputQuality != "" {
+		formats = formats.Quality(outputQuality)
 	}
-
-	var format *youtube.Format
-	itag, _ := strconv.Atoi(outputQuality)
-	switch {
-	case itag > 0:
-		// When an itag is specified, do not filter format with mime-type
-		format = video.Formats.FindByItag(itag)
-		if format == nil {
-			return nil, nil, fmt.Errorf("unable to find format with itag %d", itag)
-		}
-
-	case outputQuality != "":
-		format = formats.FindByQuality(outputQuality)
-		if format == nil {
-			return nil, nil, fmt.Errorf("unable to find format with quality %s", outputQuality)
-		}
-
-	default:
-		// select the first format
-		formats.Sort()
-		format = &formats[0]
+	if itag > 0 {
+		formats = formats.Itag(itag)
+	}
+	if formats == nil {
+		return nil, nil, fmt.Errorf("unable to find the specified format")
 	}
 
-	return video, format, nil
+	formats.Sort()
+
+	// select the first format
+	return video, &formats[0], nil
 }
diff --git a/cmd/youtubedr/info.go b/cmd/youtubedr/info.go
@@ -18,6 +18,7 @@ type VideoFormat struct {
 	VideoQuality  string
 	AudioQuality  string
 	AudioChannels int
+	Language      string
 	Size          int64
 	Bitrate       int
 	MimeType      string
@@ -73,6 +74,7 @@ var infoCmd = &cobra.Command{
 				Size:          size,
 				Bitrate:       bitrate,
 				MimeType:      format.MimeType,
+				Language:      format.LanguageDisplayName(),
 			})
 		}
 
@@ -102,6 +104,7 @@ func writeInfoOutput(w io.Writer, info *VideoInfo) {
 		"size [MB]",
 		"bitrate",
 		"MimeType",
+		"language",
 	})
 
 	for _, format := range info.Formats {
@@ -114,6 +117,7 @@ func writeInfoOutput(w io.Writer, info *VideoInfo) {
 			fmt.Sprintf("%0.1f", float64(format.Size)/1024/1024),
 			strconv.Itoa(format.Bitrate),
 			format.MimeType,
+			format.Language,
 		})
 	}
 
diff --git a/cmd/youtubedr/url.go b/cmd/youtubedr/url.go
@@ -23,7 +23,6 @@ var urlCmd = &cobra.Command{
 }
 
 func init() {
-	addQualityFlag(urlCmd.Flags())
-	addMimeTypeFlag(urlCmd.Flags())
+	addVideoSelectionFlags(urlCmd.Flags())
 	rootCmd.AddCommand(urlCmd)
 }
diff --git a/downloader/downloader.go b/downloader/downloader.go
@@ -2,7 +2,7 @@ package downloader
 
 import (
 	"context"
-	"fmt"
+	"errors"
 	"io"
 	"os"
 	"os/exec"
@@ -59,8 +59,8 @@ func (dl *Downloader) Download(ctx context.Context, v *youtube.Video, format *yo
 }
 
 // DownloadComposite : Downloads audio and video streams separately and merges them via ffmpeg.
-func (dl *Downloader) DownloadComposite(ctx context.Context, outputFile string, v *youtube.Video, quality string, mimetype string) error {
-	videoFormat, audioFormat, err1 := getVideoAudioFormats(v, quality, mimetype)
+func (dl *Downloader) DownloadComposite(ctx context.Context, outputFile string, v *youtube.Video, quality string, mimetype, language string) error {
+	videoFormat, audioFormat, err1 := getVideoAudioFormats(v, quality, mimetype, language)
 	if err1 != nil {
 		return err1
 	}
@@ -122,8 +122,7 @@ func (dl *Downloader) DownloadComposite(ctx context.Context, outputFile string,
 	return ffmpegVersionCmd.Run()
 }
 
-func getVideoAudioFormats(v *youtube.Video, quality string, mimetype string) (*youtube.Format, *youtube.Format, error) {
-	var videoFormat, audioFormat *youtube.Format
+func getVideoAudioFormats(v *youtube.Video, quality string, mimetype, language string) (*youtube.Format, *youtube.Format, error) {
 	var videoFormats, audioFormats youtube.FormatList
 
 	formats := v.Formats
@@ -138,25 +137,22 @@ func getVideoAudioFormats(v *youtube.Video, quality string, mimetype string) (*y
 		videoFormats = videoFormats.Quality(quality)
 	}
 
-	if len(videoFormats) > 0 {
-		videoFormats.Sort()
-		videoFormat = &videoFormats[0]
+	if language != "" {
+		audioFormats = audioFormats.Language(language)
 	}
 
-	if len(audioFormats) > 0 {
-		audioFormats.Sort()
-		audioFormat = &audioFormats[0]
+	if len(videoFormats) == 0 {
+		return nil, nil, errors.New("no video format found after filtering")
 	}
 
-	if videoFormat == nil {
-		return nil, nil, fmt.Errorf("no video format found after filtering")
+	if len(audioFormats) == 0 {
+		return nil, nil, errors.New("no audio format found after filtering")
 	}
 
-	if audioFormat == nil {
-		return nil, nil, fmt.Errorf("no audio format found after filtering")
-	}
+	videoFormats.Sort()
+	audioFormats.Sort()
 
-	return videoFormat, audioFormat, nil
+	return &videoFormats[0], &audioFormats[0], nil
 }
 
 func (dl *Downloader) videoDLWorker(ctx context.Context, out *os.File, video *youtube.Video, format *youtube.Format) error {
diff --git a/downloader/downloader_hq_test.go b/downloader/downloader_hq_test.go
@@ -16,6 +16,5 @@ func TestDownload_HighQuality(t *testing.T) {
 
 	video, err := testDownloader.Client.GetVideoContext(ctx, "BaW_jenozKc")
 	require.NoError(err)
-
-	require.NoError(testDownloader.DownloadComposite(ctx, "", video, "hd1080", "mp4"))
+	require.NoError(testDownloader.DownloadComposite(ctx, "", video, "hd1080", "mp4", ""))
 }
diff --git a/downloader/downloader_test.go b/downloader/downloader_test.go
@@ -69,7 +69,7 @@ func TestYoutube_DownloadWithHighQualityFails(t *testing.T) {
 				Formats: tt.formats,
 			}
 
-			err := testDownloader.DownloadComposite(context.Background(), "", video, "hd1080", "")
+			err := testDownloader.DownloadComposite(context.Background(), "", video, "hd1080", "", "")
 			assert.EqualError(t, err, tt.message)
 		})
 	}
@@ -101,7 +101,7 @@ func Test_getVideoAudioFormats(t *testing.T) {
 		{ItagNo: 249, MimeType: "audio/webm; codecs=\"opus\"", Quality: "tiny", Bitrate: 72862, FPS: 0, Width: 0, Height: 0, LastModified: "1540474783513282", ContentLength: 24839529, QualityLabel: "", ProjectionType: "RECTANGULAR", AverageBitrate: 55914, AudioQuality: "AUDIO_QUALITY_LOW", ApproxDurationMs: "3553941", AudioSampleRate: "48000", AudioChannels: 2},
 	}}
 	{
-		videoFormat, audioFormat, err := getVideoAudioFormats(v, "hd720", "mp4")
+		videoFormat, audioFormat, err := getVideoAudioFormats(v, "hd720", "mp4", "")
 		require.NoError(err)
 		require.NotNil(videoFormat)
 		require.Equal(398, videoFormat.ItagNo)
@@ -110,7 +110,7 @@ func Test_getVideoAudioFormats(t *testing.T) {
 	}
 
 	{
-		videoFormat, audioFormat, err := getVideoAudioFormats(v, "large", "webm")
+		videoFormat, audioFormat, err := getVideoAudioFormats(v, "large", "webm", "")
 		require.NoError(err)
 		require.NotNil(videoFormat)
 		require.Equal(244, videoFormat.ItagNo)
diff --git a/errors_test.go b/errors_test.go
@@ -8,6 +8,8 @@ import (
 )
 
 func TestErrors(t *testing.T) {
+	t.Parallel()
+
 	tests := []struct {
 		err      error
 		expected string
diff --git a/format_list.go b/format_list.go
@@ -8,68 +8,59 @@ import (
 
 type FormatList []Format
 
-// FindByQuality returns the first format matching Quality or QualityLabel
-//
-// Examples: tiny, small, medium, large, 720p, hd720, hd1080
-func (list FormatList) FindByQuality(quality string) *Format {
+// Type returns a new FormatList filtered by itag
+func (list FormatList) Select(f func(Format) bool) (result FormatList) {
 	for i := range list {
-		if list[i].Quality == quality || list[i].QualityLabel == quality {
-			return &list[i]
+		if f(list[i]) {
+			result = append(result, list[i])
 		}
 	}
-	return nil
+	return result
 }
 
-// FindByItag returns the first format matching the itag number
-func (list FormatList) FindByItag(itagNo int) *Format {
-	for i := range list {
-		if list[i].ItagNo == itagNo {
-			return &list[i]
-		}
-	}
-	return nil
+// Type returns a new FormatList filtered by itag
+func (list FormatList) Itag(itagNo int) FormatList {
+	return list.Select(func(f Format) bool {
+		return f.ItagNo == itagNo
+	})
 }
 
-// Type returns a new FormatList filtered by mime type of video
-func (list FormatList) Type(t string) (result FormatList) {
-	for i := range list {
-		if strings.Contains(list[i].MimeType, t) {
-			result = append(result, list[i])
-		}
-	}
-	return result
+// Type returns a new FormatList filtered by mime type
+func (list FormatList) Type(value string) FormatList {
+	return list.Select(func(f Format) bool {
+		return strings.Contains(f.MimeType, value)
+	})
+}
+
+// Type returns a new FormatList filtered by display name
+func (list FormatList) Language(displayName string) FormatList {
+	return list.Select(func(f Format) bool {
+		return f.LanguageDisplayName() == displayName
+	})
 }
 
 // Quality returns a new FormatList filtered by quality, quality label or itag,
 // but not audio quality
-func (list FormatList) Quality(quality string) (result FormatList) {
-	for _, f := range list {
-		itag, _ := strconv.Atoi(quality)
-		if itag == f.ItagNo || strings.Contains(f.Quality, quality) || strings.Contains(f.QualityLabel, quality) {
-			result = append(result, f)
-		}
-	}
-	return result
+func (list FormatList) Quality(quality string) FormatList {
+	itag, _ := strconv.Atoi(quality)
+
+	return list.Select(func(f Format) bool {
+		return itag == f.ItagNo || strings.Contains(f.Quality, quality) || strings.Contains(f.QualityLabel, quality)
+	})
 }
 
 // AudioChannels returns a new FormatList filtered by the matching AudioChannels
-func (list FormatList) AudioChannels(n int) (result FormatList) {
-	for _, f := range list {
-		if f.AudioChannels == n {
-			result = append(result, f)
-		}
-	}
-	return result
+func (list FormatList) AudioChannels(n int) FormatList {
+	return list.Select(func(f Format) bool {
+		return f.AudioChannels == n
+	})
 }
 
 // AudioChannels returns a new FormatList filtered by the matching AudioChannels
-func (list FormatList) WithAudioChannels() (result FormatList) {
-	for _, f := range list {
-		if f.AudioChannels > 0 {
-			result = append(result, f)
-		}
-	}
-	return result
+func (list FormatList) WithAudioChannels() FormatList {
+	return list.Select(func(f Format) bool {
+		return f.AudioChannels > 0
+	})
 }
 
 // FilterQuality reduces the format list to formats matching the quality
diff --git a/format_list_test.go b/format_list_test.go
diff --git a/response_data.go b/response_data.go
diff --git a/video_test.go b/video_test.go

Original file line number	Diff line number	Diff line change
`@@ -32,8 +32,7 @@ func init() {`
`32`	`32`
`33`	`33`	`downloadCmd.Flags().StringVarP(&outputFile, "filename", "o", "", "The output file, the default is genated by the video title.")`
`34`	`34`	`downloadCmd.Flags().StringVarP(&outputDir, "directory", "d", ".", "The output directory.")`
`35`		`- addQualityFlag(downloadCmd.Flags())`
`36`		`- addMimeTypeFlag(downloadCmd.Flags())`
	`35`	`+ addVideoSelectionFlags(downloadCmd.Flags())`
`37`	`36`	`}`
`38`	`37`
`39`	`38`	`func download(id string) error {`
`@@ -48,7 +47,7 @@ func download(id string) error {`
`48`	`47`	`if err := checkFFMPEG(); err != nil {`
`49`	`48`	`return err`
`50`	`49`	`}`
`51`		`- return downloader.DownloadComposite(context.Background(), outputFile, video, outputQuality, mimetype)`
	`50`	`+ return downloader.DownloadComposite(context.Background(), outputFile, video, outputQuality, mimetype, language)`
`52`	`51`	`}`
`53`	`52`
`54`	`53`	`return downloader.Download(context.Background(), video, format, outputFile)`
Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,6 @@ var urlCmd = &cobra.Command{`
`23`	`23`	`}`
`24`	`24`
`25`	`25`	`func init() {`
`26`		`- addQualityFlag(urlCmd.Flags())`
`27`		`- addMimeTypeFlag(urlCmd.Flags())`
	`26`	`+ addVideoSelectionFlags(urlCmd.Flags())`
`28`	`27`	`rootCmd.AddCommand(urlCmd)`
`29`	`28`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,5 @@ func TestDownload_HighQuality(t *testing.T) {`
`16`	`16`
`17`	`17`	`video, err := testDownloader.Client.GetVideoContext(ctx, "BaW_jenozKc")`
`18`	`18`	`require.NoError(err)`
`19`		`-`
`20`		`- require.NoError(testDownloader.DownloadComposite(ctx, "", video, "hd1080", "mp4"))`
	`19`	`+ require.NoError(testDownloader.DownloadComposite(ctx, "", video, "hd1080", "mp4", ""))`
`21`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ func TestYoutube_DownloadWithHighQualityFails(t *testing.T) {`
`69`	`69`	`Formats: tt.formats,`
`70`	`70`	`}`
`71`	`71`
`72`		`- err := testDownloader.DownloadComposite(context.Background(), "", video, "hd1080", "")`
	`72`	`+ err := testDownloader.DownloadComposite(context.Background(), "", video, "hd1080", "", "")`
`73`	`73`	`assert.EqualError(t, err, tt.message)`
`74`	`74`	`})`
`75`	`75`	`}`
`@@ -101,7 +101,7 @@ func Test_getVideoAudioFormats(t *testing.T) {`
`101`	`101`	`{ItagNo: 249, MimeType: "audio/webm; codecs=\"opus\"", Quality: "tiny", Bitrate: 72862, FPS: 0, Width: 0, Height: 0, LastModified: "1540474783513282", ContentLength: 24839529, QualityLabel: "", ProjectionType: "RECTANGULAR", AverageBitrate: 55914, AudioQuality: "AUDIO_QUALITY_LOW", ApproxDurationMs: "3553941", AudioSampleRate: "48000", AudioChannels: 2},`
`102`	`102`	`}}`
`103`	`103`	`{`
`104`		`- videoFormat, audioFormat, err := getVideoAudioFormats(v, "hd720", "mp4")`
	`104`	`+ videoFormat, audioFormat, err := getVideoAudioFormats(v, "hd720", "mp4", "")`
`105`	`105`	`require.NoError(err)`
`106`	`106`	`require.NotNil(videoFormat)`
`107`	`107`	`require.Equal(398, videoFormat.ItagNo)`
`@@ -110,7 +110,7 @@ func Test_getVideoAudioFormats(t *testing.T) {`
`110`	`110`	`}`
`111`	`111`
`112`	`112`	`{`
`113`		`- videoFormat, audioFormat, err := getVideoAudioFormats(v, "large", "webm")`
	`113`	`+ videoFormat, audioFormat, err := getVideoAudioFormats(v, "large", "webm", "")`
`114`	`114`	`require.NoError(err)`
`115`	`115`	`require.NotNil(videoFormat)`
`116`	`116`	`require.Equal(244, videoFormat.ItagNo)`
Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,8 @@ import (`
`8`	`8`	`)`
`9`	`9`
`10`	`10`	`func TestErrors(t *testing.T) {`
	`11`	`+ t.Parallel()`
	`12`	`+`
`11`	`13`	`tests := []struct {`
`12`	`14`	`err error`
`13`	`15`	`expected string`