Skip to content

Commit

Permalink
optimize code
Browse files Browse the repository at this point in the history
  • Loading branch information
Blank committed Sep 29, 2019
1 parent b5f6f3e commit 80290b8
Show file tree
Hide file tree
Showing 32 changed files with 511 additions and 242 deletions.
60 changes: 39 additions & 21 deletions downloader/downloader.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,11 @@ import (
"io"
"net/http"
"os"
"sync"
"time"

"github.com/cheggaaa/pb"

"github.com/iawia002/annie/config"
"github.com/iawia002/annie/request"
"github.com/iawia002/annie/utils"
Expand All @@ -29,7 +31,7 @@ func Caption(url, refer, fileName, ext string) error {
return nil
}
fmt.Println("\nDownloading captions...")
body, err := request.Get(url, refer, nil)
body, err := request.GetByte(url, refer, nil)
if err != nil {
return err
}
Expand All @@ -42,23 +44,27 @@ func Caption(url, refer, fileName, ext string) error {
return fileError
}
defer file.Close()
file.WriteString(body)

if _, err = file.Write(body); err != nil {
return err
}
return nil
}

func writeFile(
url string, file *os.File, headers map[string]string, bar *pb.ProgressBar,
) (int64, error) {
res, err := request.Request("GET", url, nil, headers)
res, err := request.Request(http.MethodGet, url, nil, headers)
if err != nil {
return 0, err
}
defer res.Body.Close()

writer := io.MultiWriter(file, bar)
// Note that io.Copy reads 32kb(maximum) from input and writes them to output, then repeats.
// So don't worry about memory.
written, copyErr := io.Copy(writer, res.Body)
if copyErr != nil {
if copyErr != nil && copyErr != io.EOF {
return written, fmt.Errorf("file copy error: %s", copyErr)
}
return written, nil
Expand Down Expand Up @@ -110,6 +116,17 @@ func Save(
if fileError != nil {
return fileError
}

// close and rename temp file at the end of this function
defer func() {
// must close the file before rename or it will cause
// `The process cannot access the file because it is being used by another process.` error.
file.Close()
if err == nil {
os.Rename(tempFilePath, filePath)
}
}()

if chunkSizeMB > 0 {
var start, end, chunkSize int64
chunkSize = int64(chunkSizeMB) * 1024 * 1024
Expand Down Expand Up @@ -155,15 +172,6 @@ func Save(
}
}

// close and rename temp file at the end of this function
defer func() {
// must close the file before rename or it will cause
// `The process cannot access the file because it is being used by another process.` error.
file.Close()
if err == nil {
os.Rename(tempFilePath, filePath)
}
}()
return nil
}

Expand Down Expand Up @@ -220,20 +228,25 @@ func Download(v Data, refer string, chunkSizeMB int) error {
return err
}
reqURL := fmt.Sprintf("%s://%s/jsonrpc", config.Aria2Method, config.Aria2Addr)
req, err := http.NewRequest("POST", reqURL, bytes.NewBuffer(jsonData))
req, err := http.NewRequest(http.MethodPost, reqURL, bytes.NewBuffer(jsonData))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
var client http.Client
_, err = client.Do(req)

var client = http.Client{Timeout: 30 * time.Second}
res, err := client.Do(req)
if err != nil {
return err
}
// The http Client and Transport guarantee that Body is always
// non-nil, even on responses without a body or responses with
// a zero-length body.
res.Body.Close()
}
return nil
}
var err error

// Skip the complete file that has been merged
mergedFilePath, err := utils.FilePath(title, "mp4", false)
if err != nil {
Expand Down Expand Up @@ -262,8 +275,13 @@ func Download(v Data, refer string, chunkSizeMB int) error {
wgp := utils.NewWaitGroupPool(config.ThreadNumber)
// multiple fragments
errs := make([]error, 0)
lock := sync.Mutex{}
parts := make([]string, len(data.URLs))
for index, url := range data.URLs {
if len(errs) > 0 {
break
}

partFileName := fmt.Sprintf("%s[%d]", title, index)
partFilePath, err := utils.FilePath(partFileName, url.Ext, false)
if err != nil {
Expand All @@ -276,7 +294,9 @@ func Download(v Data, refer string, chunkSizeMB int) error {
defer wgp.Done()
err := Save(url, refer, fileName, bar, chunkSizeMB)
if err != nil {
lock.Lock()
errs = append(errs, err)
lock.Unlock()
}
}(url, refer, partFileName, bar)
}
Expand All @@ -296,8 +316,6 @@ func Download(v Data, refer string, chunkSizeMB int) error {
} else {
err = utils.MergeToMP4(parts, mergedFilePath, title)
}
if err != nil {
return err
}
return nil

return err
}
3 changes: 0 additions & 3 deletions downloader/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,6 @@ type Data struct {
URL string `json:"url"`
}

// EmptyList empty Data list
var EmptyList = make([]Data, 0)

// EmptyData returns an "empty" Data object with the given URL and error
func EmptyData(url string, err error) Data {
return Data{
Expand Down
22 changes: 15 additions & 7 deletions extractors/bcy/bcy.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@ package bcy

import (
"encoding/json"
"fmt"
"strings"

"github.com/iawia002/annie/downloader"
"github.com/iawia002/annie/extractors"
"github.com/iawia002/annie/parser"
"github.com/iawia002/annie/request"
"github.com/iawia002/annie/utils"
Expand All @@ -22,21 +24,27 @@ type bcyData struct {

// Extract is the main function for extracting data
func Extract(url string) ([]downloader.Data, error) {
var err error
html, err := request.Get(url, url, nil)
if err != nil {
return downloader.EmptyList, err
return nil, err
}

// parse json data
rep := strings.NewReplacer(`\"`, `"`, `\\`, `\`)
jsonString := rep.Replace(utils.MatchOneOf(html, `JSON.parse\("(.+?)"\);`)[1])
realURLs := utils.MatchOneOf(html, `JSON.parse\("(.+?)"\);`)
if realURLs == nil || len(realURLs) < 2 {
return nil, extractors.ErrURLParseFailed
}
jsonString := rep.Replace(realURLs[1])

var data bcyData
json.Unmarshal([]byte(jsonString), &data)
if err = json.Unmarshal([]byte(jsonString), &data); err != nil {
return nil, fmt.Errorf("json unmarshal failed, err: %v", err)
}

doc, err := parser.GetDoc(html)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
title := strings.Replace(parser.Title(doc), " - 半次元 banciyuan - ACG爱好者社区", "", -1)

Expand All @@ -45,12 +53,12 @@ func Extract(url string) ([]downloader.Data, error) {
for _, img := range data.Detail.PostData.Multi {
size, err := request.Size(img.OriginalPath, url)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
totalSize += size
_, ext, err := utils.GetNameAndExt(img.OriginalPath)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
urls = append(urls, downloader.URL{
URL: img.OriginalPath,
Expand Down
16 changes: 12 additions & 4 deletions extractors/bilibili/bilibili.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (

"github.com/iawia002/annie/config"
"github.com/iawia002/annie/downloader"
"github.com/iawia002/annie/extractors"
"github.com/iawia002/annie/parser"
"github.com/iawia002/annie/request"
"github.com/iawia002/annie/utils"
Expand Down Expand Up @@ -111,7 +112,7 @@ func extractBangumi(url, html string) ([]downloader.Data, error) {
var data bangumiData
err := json.Unmarshal([]byte(dataString), &data)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
if !config.Playlist {
options := bilibiliOptions{
Expand Down Expand Up @@ -173,7 +174,7 @@ func getMultiPageData(html string) (*multiPage, error) {
func extractNormalVideo(url, html string) ([]downloader.Data, error) {
pageData, err := getMultiPageData(html)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
if !config.Playlist {
// handle URL that has a playlist, mainly for unified titles
Expand All @@ -189,6 +190,10 @@ func extractNormalVideo(url, html string) ([]downloader.Data, error) {
p, _ = strconv.Atoi(pageString[1])
}

if len(pageData.VideoData.Pages) < p || p < 1 {
return nil, extractors.ErrURLParseFailed
}

page := pageData.VideoData.Pages[p-1]
options := bilibiliOptions{
url: url,
Expand Down Expand Up @@ -240,7 +245,7 @@ func Extract(url string) ([]downloader.Data, error) {
var err error
html, err := request.Get(url, referer, nil)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
if strings.Contains(url, "bangumi") {
// handle bangumi
Expand Down Expand Up @@ -330,10 +335,13 @@ func bilibiliDownload(options bilibiliOptions) downloader.Data {
title = tempTitle
}

downloader.Caption(
err = downloader.Caption(
fmt.Sprintf("https://comment.bilibili.com/%s.xml", options.cid),
options.url, title, "xml",
)
if err != nil {
return downloader.EmptyData(options.url, err)
}

return downloader.Data{
Site: "哔哩哔哩 bilibili.com",
Expand Down
7 changes: 7 additions & 0 deletions extractors/defs.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
package extractors

import (
"errors"
)

var ErrURLParseFailed = errors.New("url parse failed")
12 changes: 9 additions & 3 deletions extractors/douyin/douyin.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package douyin

import (
"github.com/iawia002/annie/downloader"
"github.com/iawia002/annie/extractors"
"github.com/iawia002/annie/request"
"github.com/iawia002/annie/utils"
)
Expand All @@ -11,7 +12,7 @@ func Extract(url string) ([]downloader.Data, error) {
var err error
html, err := request.Get(url, url, nil)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
var title string
desc := utils.MatchOneOf(html, `<p class="desc">(.+?)</p>`)
Expand All @@ -20,10 +21,15 @@ func Extract(url string) ([]downloader.Data, error) {
} else {
title = "抖音短视频"
}
realURL := utils.MatchOneOf(html, `playAddr: "(.+?)"`)[1]
realURLs := utils.MatchOneOf(html, `playAddr: "(.+?)"`)
if realURLs == nil || len(realURLs) < 2 {
return nil, extractors.ErrURLParseFailed
}
realURL := realURLs[1]

size, err := request.Size(realURL, url)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
urlData := downloader.URL{
URL: realURL,
Expand Down
22 changes: 16 additions & 6 deletions extractors/douyu/douyu.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"errors"

"github.com/iawia002/annie/downloader"
"github.com/iawia002/annie/extractors"
"github.com/iawia002/annie/request"
"github.com/iawia002/annie/utils"
)
Expand Down Expand Up @@ -52,26 +53,35 @@ func Extract(url string) ([]downloader.Data, error) {
var err error
liveVid := utils.MatchOneOf(url, `https?://www.douyu.com/(\S+)`)
if liveVid != nil {
return downloader.EmptyList, errors.New("暂不支持斗鱼直播")
return nil, errors.New("暂不支持斗鱼直播")
}

html, err := request.Get(url, url, nil)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
title := utils.MatchOneOf(html, `<title>(.*?)</title>`)[1]
titles := utils.MatchOneOf(html, `<title>(.*?)</title>`)
if titles == nil || len(titles) < 2 {
return nil, extractors.ErrURLParseFailed
}
title := titles[1]

vids := utils.MatchOneOf(url, `https?://v.douyu.com/show/(\S+)`)
if vids == nil || len(vids) < 2 {
return nil, extractors.ErrURLParseFailed
}
vid := vids[1]

vid := utils.MatchOneOf(url, `https?://v.douyu.com/show/(\S+)`)[1]
dataString, err := request.Get("http://vmobile.douyu.com/video/getInfo?vid="+vid, url, nil)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
var dataDict douyuData
json.Unmarshal([]byte(dataString), &dataDict)

m3u8URLs, totalSize, err := douyuM3u8(dataDict.Data.VideoURL)
if err != nil {
return downloader.EmptyList, err
return nil, err
}
urls := make([]downloader.URL, len(m3u8URLs))
for index, u := range m3u8URLs {
Expand Down
Loading

0 comments on commit 80290b8

Please sign in to comment.