Add support for more URL formats

closes ä60
rylio · Dec 27, 2019 · 5ae23c9 · 5ae23c9
1 parent c888425
commit 5ae23c9
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 20 deletions.
diff --git a/video_info.go b/video_info.go
@@ -4,7 +4,6 @@ import (
  "bytes"
  "encoding/json"
  "encoding/xml"
- "errors"
  "fmt"
  "io"
  "net/http"
@@ -46,46 +45,50 @@ func GetVideoInfo(value interface{}) (*VideoInfo, error) {
  case *url.URL:
  return GetVideoInfoFromURL(t)
  case string:
- u, err := url.ParseRequestURI(t)
- if err != nil {
- return GetVideoInfoFromID(t)
- }
- if u.Host == "youtu.be" {
- return GetVideoInfoFromShortURL(u)
+ if strings.HasPrefix(t, "https://") {
+  uri, err := url.ParseRequestURI(t)
+ if err != nil {
+  return nil, err
+  }
+ return GetVideoInfo(uri)
  }
- return GetVideoInfoFromURL(u)
+
+ return GetVideoInfoFromID(t)
  default:
  return nil, fmt.Errorf("Identifier type must be a string, *url.URL, or []byte")
  }
 }
 
 // GetVideoInfoFromURL fetches video info from a youtube url
 func GetVideoInfoFromURL(u *url.URL) (*VideoInfo, error) {
- videoID := u.Query().Get("v")
+ videoID := extractVideoID(u)
  if len(videoID) == 0 {
  return nil, fmt.Errorf("Invalid youtube url, no video id")
  }
  return GetVideoInfoFromID(videoID)
 }
 
 // GetVideoInfoFromShortURL fetches video info from a short youtube url
-func GetVideoInfoFromShortURL(u *url.URL) (*VideoInfo, error) {
- if len(u.Path) >= 1 {
- if path := u.Path[1:]; path != "" {
- return GetVideoInfoFromID(path)
+func extractVideoID(u *url.URL) string {
+ switch u.Host {
+ case "www.youtube.com", "youtube.com":
+ if u.Path == "/watch" {
+ return u.Query().Get("v")
+ }
+ if strings.HasPrefix(u.Path, "/embed/") {
+ return u.Path[7:]
+ }
+ case "youtu.be":
+ if len(u.Path) > 1 {
+ return u.Path[1:]
  }
  }
- return nil, errors.New("Could not parse short URL")
+ return ""
 }
 
 // GetVideoInfoFromID fetches video info from a youtube video id
 func GetVideoInfoFromID(id string) (*VideoInfo, error) {
- u, _ := url.ParseRequestURI(youtubeBaseURL)
- values := u.Query()
- values.Set("v", id)
- u.RawQuery = values.Encode()
-
- body, err := httpGetAndCheckResponseReadBody(u.String())
+ body, err := httpGetAndCheckResponseReadBody(youtubeBaseURL + "?v=" + id)
 
  if err != nil {
  return nil, err

diff --git a/video_info_test.go b/video_info_test.go
@@ -3,6 +3,7 @@ package ytdl
 import (
  "io/ioutil"
  "net/http"
+ "net/url"
  "testing"
  "time"
 
@@ -122,6 +123,29 @@ Linus gives the practical reasons why he doesn't use Ubuntu or Debian.`,
  }
 }
 
+func TestExtractIDfromValidURL(t *testing.T) {
+ tests := []string{
+ "https://youtube.com/watch?v=BaW_jenozKc",
+ "https://www.youtube.com/watch?v=BaW_jenozKc",
+ "https://www.youtube.com/watch?v=BaW_jenozKc&v=UxxajLWwzqY",
+ "https://www.youtube.com/embed/BaW_jenozKc?list=PLEbnTDJUr_IegfoqO4iPnPYQui46QqT0j",
+ "https://youtu.be/BaW_jenozKc",
+ }
+ for _, input := range tests {
+ t.Run(input, func(t *testing.T) {
+ uri, err := url.ParseRequestURI(input)
+ require.NoError(t, err)
+ assert.Equal(t, "BaW_jenozKc", extractVideoID(uri))
+ })
+ }
+}
+
+func TestExtractIDfromInvalidURL(t *testing.T) {
+ uri, err := url.ParseRequestURI("https://otherhost.com/watch?v=BaW_jenozKc")
+ require.NoError(t, err)
+ assert.Equal(t, "", extractVideoID(uri))
+}
+
 func TestGetDownloadURL(t *testing.T) {
  testCases := []string{
  "FrG4TEcSuRg",