extractors/pornhub: support different quality

README.md

@@ -600,6 +600,7 @@ Twitter | <https://twitter.com> | ✓ | | | |
 网易云音乐 | <https://music.163.com> | ✓ | | | |
 音悦台 | <https://yinyuetai.com> | ✓ | | | |
 极客时间 | <https://time.geekbang.org> | ✓ | | | |
+Pornhub | <https://pornhub.com> | ✓ | | | |
 
 
 ## Known issues

extractors/pornhub/pornhub.go

@@ -2,17 +2,17 @@ package pornhub
 
 import (
 	"encoding/json"
+	"fmt"
 
 	"github.com/iawia002/annie/downloader"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
 )
 
 type pornhubData struct {
-	DefaultQuality bool   `json:"defaultQuality"`
-	Format         string `json:"format"`
-	Quality        string `json:"quality"`
-	VideoURL       string `json:"videoUrl"`
+	Format   string `json:"format"`
+	Quality  string `json:"quality"`
+	VideoURL string `json:"videoUrl"`
 }
 
 // Extract is the main function for extracting data
@@ -38,37 +38,35 @@ func Extract(url string) ([]downloader.Data, error) {
 		return downloader.EmptyList, err
 	}
 
-	//TODO add support for different quality
-	var realURL string
-	for _, downloadlink := range(pornhubs) {
-		if downloadlink.VideoURL != "" {
-			realURL = downloadlink.VideoURL
-			break
+	streams := make(map[string]downloader.Stream, len(pornhubs))
+	for _, data := range pornhubs {
+		realURL := data.VideoURL
+		if realURL == "" {
+			continue
 		}
-	}
-
-	size, err := request.Size(realURL, url)
-	if err != nil {
-		return downloader.EmptyList, err
-	}
-	urlData := downloader.URL{
-		URL:  realURL,
-		Size: size,
-		Ext:  "mp4",
-	}
-	streams := map[string]downloader.Stream{
-		"default": {
-			URLs: []downloader.URL{urlData},
+		size, err := request.Size(realURL, url)
+		if err != nil {
+			return downloader.EmptyList, err
+		}
+		urlData := downloader.URL{
+			URL:  realURL,
 			Size: size,
-		},
+			Ext:  "mp4",
+		}
+		streams[data.Quality] = downloader.Stream{
+			URLs:    []downloader.URL{urlData},
+			Size:    size,
+			Quality: fmt.Sprintf("%sP", data.Quality),
+		}
 	}
+
 	return []downloader.Data{
 		{
-			Site:    "Pornhub",
+			Site:    "Pornhub pornhub.com",
 			Title:   title,
 			Type:    "video",
 			Streams: streams,
 			URL:     url,
 		},
 	}, nil
-}
+}

extractors/pornhub/pornhub_test.go

@@ -7,8 +7,7 @@ import (
 	"github.com/iawia002/annie/test"
 )
 
-
-func TestYoutube(t *testing.T) {
+func TestPornhub(t *testing.T) {
 	config.InfoOnly = true
 	config.RetryTimes = 10
 	tests := []struct {
@@ -18,9 +17,10 @@ func TestYoutube(t *testing.T) {
 		{
 			name: "normal test",
 			args: test.Args{
-				URL:   "https://www.pornhub.com/view_video.php?viewkey=ph5cb5fc41c6ebd",
-				Title: "Must Watch MILF Drilled by the Fireplace",
-				Size: 158868371,
+				URL:     "https://www.pornhub.com/view_video.php?viewkey=ph5cb5fc41c6ebd",
+				Title:   "Must watch Milf drilled by the fireplace",
+				Quality: "720P",
+				Size:    158868371,
 			},
 		},
 	}
@@ -31,4 +31,4 @@ func TestYoutube(t *testing.T) {
 			test.Check(t, tt.args, data[0])
 		})
 	}
-}
+}