Add http support and format code

Signed-off-by: wujiawei <wujiawei@xiaomi.com>
einverne · Aug 3, 2017 · fe58e98 · fe58e98
1 parent 44cde5f
commit fe58e98
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 15 deletions.
diff --git a/douban/__main__.py b/douban/__main__.py
@@ -32,19 +32,19 @@ def parse_url(url, path):
     :param url:
     :return:
     """
-    match = re.match(r'https://www.douban.com/photos/album/(\d+)', url)
+    match = re.match(r'(http|https)://www.douban.com/photos/album/(\d+)', url)
     if match:
-        album_id = match.group(1)
+        album_id = match.group(2)
         album = Album(album_id)
         get_album(album, path)
         return
-    match = re.match(r'https://movie.douban.com/celebrity/(\d+)', url)
+    match = re.match(r'(http|https)://movie.douban.com/celebrity/(\d+)', url)
     if match:
-        celebrity_id = match.group(1)
+        celebrity_id = match.group(2)
         celebrity = Celebrity(celebrity_id)
         get_celebrity(celebrity, path)
         return
-    print "Not support this url yet"
+    print("Not support this url yet")
 
 
 def main():

diff --git a/douban/douban_album_dl.py b/douban/douban_album_dl.py
@@ -2,8 +2,6 @@
 
 import os
 
-import requests
-
 from utils import file_utils
 
 headers = {
@@ -22,8 +20,6 @@ def get_album(album, path):
     for photo_url in album.photos():
         name = os.path.basename(photo_url)
         print("{}: saving {}".format(idx, name))
-        r = requests.get(photo_url, headers=headers, stream=True)
-        with open(name, "wb") as f:
-            f.write(r.content)
+        file_utils.save_from_url(photo_url, headers, name)
         idx += 1
     print("saving album to {}, total {} images".format(path, idx))
diff --git a/douban/douban_celebrity_dl.py b/douban/douban_celebrity_dl.py
@@ -2,9 +2,8 @@
 # -*- coding:utf-8 -*-
 import os
 
-import requests
-
 from douban.celebrity import Celebrity
+from utils import file_utils
 
 
 def get_celebrity(celebrity, path):
@@ -22,9 +21,7 @@ def get_celebrity(celebrity, path):
             "Referer": refer,
             "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/59.0.3071.104 Safari/537.36"
         }
-        r = requests.get(photo_url, headers=headers, stream=True)
-        with open(name, "wb") as f:
-            f.write(r.content)
+        file_utils.save_from_url(photo_url, headers, name)
         idx += 1
     print("saving celebrity photo to {}".format(path))
 

diff --git a/utils/file_utils.py b/utils/file_utils.py
@@ -2,7 +2,15 @@
 # -*- coding: utf-8 -*-
 import os
 
+import requests
+
 
 def mkdir(path):
     if not os.path.exists(path):
         os.makedirs(path)
+
+
+def save_from_url(url, headers, name):
+    r = requests.get(url, headers=headers, stream=True)
+    with open(name, 'wb') as f:
+        f.write(r.content)