forked from andeya/faygo
/
util.go
133 lines (115 loc) · 2.93 KB
/
util.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
// Copyright 2015 henrylee2cn Author. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package surfer
import (
"io"
"io/ioutil"
"log"
"net/http"
"net/url"
"os"
"path/filepath"
"strings"
"golang.org/x/net/html/charset"
)
// AutoToUTF8 采用surf内核下载时,可以尝试自动转码为utf8
// 采用phantomjs内核时,无需转码(已是utf8)
func AutoToUTF8(resp *http.Response) error {
destReader, err := charset.NewReader(resp.Body, resp.Header.Get("Content-Type"))
if err == nil {
resp.Body = &RespBody{
ReadCloser: resp.Body,
Reader: destReader,
}
}
return err
}
// BodyBytes 读取完整响应流正文
func BodyBytes(resp *http.Response) ([]byte, error) {
body, err := ioutil.ReadAll(resp.Body)
resp.Body.Close()
return body, err
}
// UrlEncode 返回编码后的url.URL指针、及解析错误
func UrlEncode(urlStr string) (*url.URL, error) {
urlObj, err := url.Parse(urlStr)
urlObj.RawQuery = urlObj.Query().Encode()
return urlObj, err
}
// GetWDPath gets the work directory path.
func GetWDPath() string {
wd := os.Getenv("GOPATH")
if wd == "" {
panic("GOPATH is not setted in env.")
}
return wd
}
// IsDirExists judges path is directory or not.
func IsDirExists(path string) bool {
fi, err := os.Stat(path)
if err != nil {
return os.IsExist(err)
} else {
return fi.IsDir()
}
panic("util isDirExists not reached")
}
// IsFileExists judges path is file or not.
func IsFileExists(path string) bool {
fi, err := os.Stat(path)
if err != nil {
return os.IsExist(err)
} else {
return !fi.IsDir()
}
panic("util isFileExists not reached")
}
// WalkDir 遍历目录,可指定后缀
func WalkDir(targpath string, suffixes ...string) (dirlist []string) {
if !filepath.IsAbs(targpath) {
targpath, _ = filepath.Abs(targpath)
}
err := filepath.Walk(targpath, func(retpath string, f os.FileInfo, err error) error {
if err != nil {
return err
}
if !f.IsDir() {
return nil
}
if len(suffixes) == 0 {
dirlist = append(dirlist, retpath)
return nil
}
for _, suffix := range suffixes {
if strings.HasSuffix(retpath, suffix) {
dirlist = append(dirlist, retpath)
}
}
return nil
})
if err != nil {
log.Printf("utils.WalkDir: %v\n", err)
return
}
return
}
// RespBody 封装Response.Body
type RespBody struct {
io.ReadCloser
io.Reader
}
// Read 实现Reader接口
func (b *RespBody) Read(p []byte) (int, error) {
return b.Reader.Read(p)
}