Add support for more URL formats

closes ä60

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

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",