完善上传进度

为 sm.ms 和 imgtg 添加上传进度支持。

Author: thep0y

pyproject.toml

@@ -14,7 +14,7 @@ keywords = ["typora", "image bed", "upload"]
 dependencies = [
   "requests",
   "click",
-  "colorful-logger>=0.2.0b4",
+  "colorful-logger>=0.2.0b6",
   "typing_extensions; python_version<'3.8'",
   "tqdm",
   "requests-toolbelt",

up2b/__init__.py

@@ -103,7 +103,7 @@ def login(username: str, password: str):
             "you have chosen `github` as the image bed, please login with `-lg`"
         )
 
-    logger.info("current image bed", name=ib)
+    logger.info("current image bed", name=str(ib))
 
     echo("正在验证账号，请耐心等待...")
 
@@ -177,19 +177,29 @@ def login_git(access_token: str, username: str, repository: str, path: str):
     default=False,
     help="忽略数据库缓存，强制上传图片",
 )
+@click.option(
+    "-q",
+    "--quiet",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="静默模式。开启后不显示上传进度",
+)
 @click.option("-t", "--timeout", type=float, help="上传图片的超时时间")
 def upload(
     image_paths: Tuple[str],
     add_watermark: bool,
     auto_compress: bool,
     ignore_cache: bool,
     timeout: float,
+    quiet: bool,
 ):
     ib = _read_image_bed(
         add_watermark=add_watermark,
         auto_compress=auto_compress,
         ignore_cache=ignore_cache,
         timeout=timeout,
+        quiet=quiet,
     )
 
     paths = check_paths(image_paths)
@@ -269,11 +279,12 @@ def _read_image_bed(
     add_watermark: bool = False,
     ignore_cache: bool = False,
     timeout: Optional[float] = None,
+    quiet: bool = False,
 ) -> Union[SM, Imgtu, Imgtg, Github]:
     conf = read_conf()
 
     selected_code = conf.image_bed
-    if not selected_code:
+    if selected_code == None:
         logger.fatal("当前图床为空，请先选择要使用的图床")
 
     assert isinstance(selected_code, int)
@@ -286,13 +297,15 @@ def _read_image_bed(
             add_watermark=add_watermark,
             ignore_cache=ignore_cache,
             timeout=timeout,
+            quiet=quiet,
         )
     except ValueError:
         IMAGE_BEDS[ImageBedCode.SM_MS](
             auto_compress=auto_compress,
             add_watermark=add_watermark,
             ignore_cache=ignore_cache,
             timeout=timeout,
+            quiet=quiet,
         )
         logger.fatal("未知的图床代码，可能因为清除无效的 gitee 配置，请重试", code=selected_code)
 

up2b/up2b_lib/file.py

@@ -0,0 +1,39 @@
+import mimetypes
+
+from typing import Optional, Tuple, BinaryIO, Union
+from pathlib import Path
+from up2b.up2b_lib.custom_types import ImageType
+
+
+class File:
+    def __init__(
+        self, key: str, image: ImageType, filename: Optional[str] = None
+    ) -> None:
+        self.__key = key
+
+        self.mime_type = (
+            mimetypes.guess_type(image)[0]
+            if isinstance(image, Path)
+            else image.mime_type
+        )
+
+        self.filename = (
+            filename or image.name if isinstance(image, Path) else image.filename
+        )
+        self.image = image
+
+    @property
+    def key(self):
+        return self.__key
+
+    def to_tuple(self) -> Tuple[str, Union[BinaryIO, bytes], str]:
+        return (
+            self.filename,
+            self.image.open("rb")
+            if isinstance(self.image, Path)
+            else self.image.stream,
+            self.mime_type,  # type: ignore
+        )
+
+    def to_dict(self):
+        return {self.key: self.to_tuple()}

up2b/up2b_lib/http.py

@@ -7,6 +7,7 @@
 
 from typing import Any, Dict, Optional
 from tqdm import tqdm
+from up2b.up2b_lib.file import File
 
 
 class ProgressBar(tqdm):
@@ -16,17 +17,21 @@ def update_to(self, n: int) -> None:
 
 def upload_with_progress_bar(
     url: str,
-    filename: str,
-    form: Dict[str, Any],
+    file: File,
+    timeout: float,
+    form: Optional[Dict[str, Any]] = None,
     headers: Optional[Dict[str, str]] = None,
 ):
-    encoder = requests_toolbelt.MultipartEncoder(form)
+    data = form or {}
+    data.update(file.to_dict())
+
+    encoder = requests_toolbelt.MultipartEncoder(data)
 
     headers = headers or {}
 
     with ProgressBar(
         total=encoder.len,
-        desc=filename,
+        desc=file.filename,
         unit="B",
         unit_scale=True,
         unit_divisor=1024,
@@ -39,6 +44,6 @@ def upload_with_progress_bar(
 
         headers.update({"Content-Type": monitor.content_type})
 
-        resp = requests.post(url, data=monitor, headers=headers)
+        resp = requests.post(url, data=monitor, headers=headers, timeout=timeout)
 
         return resp

up2b/up2b_lib/up2b_api/__init__.py

@@ -358,8 +358,12 @@ def __init__(
         add_watermark: bool = False,
         ignore_cache: bool = False,
         conf: Optional[Config] = None,
+        timeout: Optional[float] = None,
+        quiet: bool = False,
     ):
-        super().__init__(auto_compress, add_watermark, ignore_cache, conf)
+        super().__init__(
+            auto_compress, add_watermark, ignore_cache, conf, timeout, quiet
+        )
 
         if self.auth_info:
             self.token = self.auth_info["token"]

up2b/up2b_lib/up2b_api/__init__.pyi

@@ -1,4 +1,3 @@
-from collections.abc import Sequence
 from pathlib import Path
 import requests
 
@@ -89,6 +88,7 @@ class Base:
         ignore_cache: bool = ...,
         conf: Optional[Config] = ...,
         timeout: Optional[float] = ...,
+        quiet: bool = ...,
     ) -> None: ...
     def check_login(self) -> None: ...
     def _read_auth_info(self) -> Optional[AuthInfo]: ...
@@ -119,6 +119,7 @@ class GitBase(Base, ImageBedAbstract):
         ignore_cache: bool = ...,
         conf: Optional[Config] = ...,
         timeout: Optional[float] = None,
+        quiet: bool = ...,
     ) -> None: ...
     def login(
         self, token: str, username: str, repo: str, folder: str = ...

up2b/up2b_lib/up2b_api/github.py

@@ -1,11 +1,5 @@
 #!/usr/bin/env python3
 # -*- coding:utf-8 -*-
-# @Author:      thepoy
-# @Email:       thepoy@163.com
-# @File Name:   github.py
-# @Created At:  2021-02-13 09:10:14
-# @Modified At: 2023-03-04 21:40:12
-# @Modified By: thepoy
 
 import os
 import requests
@@ -39,8 +33,11 @@ def __init__(
         ignore_cache: bool = False,
         conf: Optional[Config] = None,
         timeout: Optional[float] = None,
+        quiet: bool = False,
     ):
-        super().__init__(auto_compress, add_watermark, ignore_cache, conf, timeout)
+        super().__init__(
+            auto_compress, add_watermark, ignore_cache, conf, timeout, quiet
+        )
 
         if hasattr(self, "token"):
             self.headers = {

up2b/up2b_lib/up2b_api/imgtg.py

@@ -5,7 +5,6 @@
 import re
 import time
 import json
-import mimetypes
 import requests
 
 from urllib import parse
@@ -22,6 +21,8 @@
     UploadErrorResponse,
 )
 from up2b.up2b_lib.errors import MissingAuth
+from up2b.up2b_lib.file import File
+from up2b.up2b_lib.http import upload_with_progress_bar
 from up2b.up2b_lib.up2b_api import Base
 from up2b.up2b_lib.log import child_logger
 from up2b.up2b_lib.constants import IMAGE_BEDS_NAME, ImageBedCode
@@ -46,8 +47,11 @@ def __init__(
         ignore_cache: bool = False,
         conf: Optional[Config] = None,
         timeout: Optional[float] = None,
+        quiet: bool = False,
     ):
-        super().__init__(auto_compress, add_watermark, ignore_cache, conf, timeout)
+        super().__init__(
+            auto_compress, add_watermark, ignore_cache, conf, timeout, quiet
+        )
 
         self.cookie: Optional[str] = None
         self.token: Optional[str] = None
@@ -172,20 +176,10 @@ def __upload(
         else:
             filename = filename_with_suffix + "." + image.mime_type
 
-        mime_type = (
-            mimetypes.guess_type(image)[0]
-            if isinstance(image, Path)
-            else image.mime_type
-        )
+        file = File("source", image, filename=filename)
 
         timestamp = int(time.time() * 1000)
 
-        if isinstance(image, Path):
-            with open(image, "rb") as fb:
-                img_buffer = fb.read()
-        else:
-            img_buffer = image.stream
-
         data = {
             "type": "file",
             "action": "upload",
@@ -194,15 +188,24 @@ def __upload(
             "nsfw": "0",
         }
 
-        files = {
-            "source": (filename, img_buffer, mime_type),
-        }
-
         logger.debug("请求头", header=self.headers)
 
-        resp = requests.post(
-            url, headers=self.headers, data=data, files=files, timeout=self.timeout  # type: ignore
-        )
+        try:
+            if not self.quiet:
+                resp = upload_with_progress_bar(
+                    url, file, self.timeout, data, self.headers
+                )
+            else:
+                resp = requests.post(
+                    url,
+                    headers=self.headers,
+                    data=data,
+                    files=file.to_dict(),
+                    timeout=self.timeout,
+                )
+        except requests.exceptions.ConnectionError as e:
+            return UploadErrorResponse(400, str(e), str(image))
+
         resp.encoding = "utf-8"
 
         logger.debug("实际请求头", header=resp.request.headers)

up2b/up2b_lib/up2b_api/imgtu.py

@@ -1,17 +1,10 @@
 #!/usr/bin/env python3
 # -*- coding:utf-8 -*-
-# @Author:      thepoy
-# @Email:       thepoy@163.com
-# @File Name:   imgtu.py
-# @Created At:  2021-02-13 09:04:37
-# @Modified At: 2023-04-19 14:48:44
-# @Modified By: thepoy
 
 import os
 import re
 import time
 import json
-import mimetypes
 import requests
 
 from urllib import parse
@@ -28,6 +21,7 @@
     UploadErrorResponse,
 )
 from up2b.up2b_lib.errors import MissingAuth
+from up2b.up2b_lib.file import File
 from up2b.up2b_lib.http import upload_with_progress_bar
 from up2b.up2b_lib.up2b_api import Base
 from up2b.up2b_lib.constants import IMAGE_BEDS_NAME, ImageBedCode
@@ -53,8 +47,11 @@ def __init__(
         ignore_cache: bool = False,
         conf: Optional[Config] = None,
         timeout: Optional[float] = None,
+        quiet: bool = False,
     ):
-        super().__init__(auto_compress, add_watermark, ignore_cache, conf, timeout)
+        super().__init__(
+            auto_compress, add_watermark, ignore_cache, conf, timeout, quiet
+        )
 
         self.cookie: Optional[str] = None
         self.token: Optional[str] = None
@@ -172,20 +169,10 @@ def __upload(
         else:
             filename = filename_with_suffix + "." + image.mime_type
 
-        mime_type = (
-            mimetypes.guess_type(image)[0]
-            if isinstance(image, Path)
-            else image.mime_type
-        )
+        file = File("source", image, filename=filename)
 
         timestamp = int(time.time() * 1000)
 
-        if isinstance(image, Path):
-            with open(image, "rb") as fb:
-                img_buffer = fb.read()
-        else:
-            img_buffer = image.stream
-
         data = {
             "type": "file",
             "action": "upload",
@@ -194,18 +181,18 @@ def __upload(
             "nsfw": "0",
         }
 
-        files = {
-            "source": (filename, img_buffer, mime_type),
-        }
-
         try:
             if not self.quiet:
                 resp = upload_with_progress_bar(
-                    url, filename, dict(data, **files), self.headers
+                    url, file, self.timeout, data, self.headers
                 )
             else:
                 resp = requests.post(
-                    url, headers=self.headers, data=data, files=files, timeout=self.timeout  # type: ignore
+                    url,
+                    headers=self.headers,
+                    data=data,
+                    files=file.to_dict(),
+                    timeout=self.timeout,
                 )
         except requests.exceptions.ConnectionError as e:
             return UploadErrorResponse(400, str(e), str(image))