Merge pull request #20 from ethpch/master

handle auto_ext in download correctly

pixivpy_async/bapi.py

@@ -1,7 +1,6 @@
 # -*- coding:utf-8 -*-
 import hashlib
 import os
-import re
 from datetime import datetime
 import aiofiles
 from .error import *
@@ -140,21 +139,24 @@ async def download(self, url, prefix='', path=os.path.curdir, fname=None, auto_e
             name = fname
         if name:
             img_path = os.path.join(path, prefix + name)
+            if auto_ext is True:
+                _ = await self.down(url, referer, _request_content_type=auto_ext)
+                type = await _.__anext__()
+                if type in self.content_type_mapping.keys():
+                    _name, _ext = os.path.splitext(img_path)
+                    img_path = _name + self.content_type_mapping[type]
             if os.path.exists(img_path) and not replace:
                 return False
             else:
-                response, type = await self.down(url, referer)
-                if auto_ext and type in self.content_type_mapping:
-                    _ext = re.findall(r'(\.\w+)$', img_path)
-                    if _ext:
-                        img_path = img_path.replace(
-                            _ext[0], self.content_type_mapping[type])
-                    else:
-                        img_path += self.content_type_mapping[type]
+                _ = locals().get('_', await self.down(url, referer, _request_content_type=False))
+                response = await _.__anext__()
+                await _.aclose()
                 async with aiofiles.open(img_path, mode='wb') as out_file:
                     await out_file.write(response)
         else:
-            response, _ = await self.down(url, referer)
+            _ = await self.down(url, referer, _request_content_type=False)
+            response = await _.__anext__()
+            await _.aclose()
             fname.write(response)
         del response
         return True

pixivpy_async/client.py

@@ -54,6 +54,7 @@ def __init__(self, limit=30, timeout=10, env=False, internal=False, proxy=None,
 
         if proxy and _flag:
             from functools import partial
+            self.client.head = partial(self.client.head, proxy=proxy)
             self.client.get = partial(self.client.get, proxy=proxy)
             self.client.post = partial(self.client.post, proxy=proxy)
 

pixivpy_async/net.py

@@ -34,14 +34,13 @@ def __init__(self, client=None, env=False, proxy=None, bypass=False, **requests_
         self.requests_kwargs = requests_kwargs
 
     @retry(*retryable_error, retries=10, cooldown=random.randint(1, 3))
-    async def down(self, _url, _referer):
+    async def down(self, _url, _referer, _request_content_type: bool = False):
         async with ClientManager(self.session, **self.conn_opt) as session:
+            if _request_content_type is True:
+                async with session.head(_url, headers={'Referer': _referer}, **self.requests_kwargs) as head:
+                    yield head.content_type
             async with session.get(_url, headers={'Referer': _referer}, **self.requests_kwargs) as res:
-                c = await res.read()
-                t = res.content_type
-
-        await asyncio.sleep(0)
-        return c, t
+                yield await res.read()
 
     @retry(*retryable_error, retries=10, cooldown=random.randint(1, 3))
     async def auth(self, _url, _headers, _data):

pixivpy_async/retry.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from inspect import isasyncgen
 
 from functools import wraps
 from .error import RetryExhaustedError
@@ -25,7 +26,9 @@ async def inner(*args, **kwargs):
 
             while True:
                 try:
-                    result = await func(*args, **kwargs)
+                    result = func(*args, **kwargs)
+                    if isasyncgen(result) is False:
+                        result = await result
                 except exceptions as err:
                     retries_count += 1
                     message = "Exception during {} execution. " \