use httpx instead of requests
httpx is already typed, upload monitoring isn't so great
This commit is contained in:
parent
a79585f52f
commit
2254ee3ec4
|
@ -16,14 +16,16 @@
|
|||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from io import BufferedIOBase, BufferedReader
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
from typing import Literal, Optional
|
||||
|
||||
import httpx
|
||||
import pyperclip
|
||||
import requests
|
||||
import typer
|
||||
from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
|
||||
from rich.console import Console
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
|
@ -40,6 +42,11 @@ HIKARI_URL = f"{HIKARI_BASE}/upload"
|
|||
console = Console()
|
||||
|
||||
|
||||
@cache
|
||||
def get_client() -> httpx.Client:
|
||||
return httpx.Client()
|
||||
|
||||
|
||||
@dataclass
|
||||
class UploadResponse:
|
||||
status: Literal["exists"] | Literal["uploaded"]
|
||||
|
@ -58,12 +65,12 @@ def check_hash(filename: str, file: Path):
|
|||
with file.open('br') as f:
|
||||
files = {'file': (filename, f.read(2048))}
|
||||
|
||||
with requests.post(url, files=files, stream=True) as req:
|
||||
if req.status_code == 404:
|
||||
return
|
||||
req = get_client().post(url, files=files)
|
||||
if req.status_code == 404:
|
||||
return
|
||||
|
||||
req.raise_for_status()
|
||||
return UploadResponse(**req.json())
|
||||
req.raise_for_status()
|
||||
return UploadResponse(**req.json())
|
||||
|
||||
|
||||
def output(copy: bool, resp: UploadResponse):
|
||||
|
@ -82,10 +89,13 @@ def output(copy: bool, resp: UploadResponse):
|
|||
|
||||
@dataclass
|
||||
class ProgressMonitor:
|
||||
size: int
|
||||
filename: str
|
||||
file: BufferedReader
|
||||
|
||||
def __post_init__(self):
|
||||
def __post_init__(self) -> None:
|
||||
self.size = get_file_size(self.file)
|
||||
self._file_read = self.file.read
|
||||
self.file.read = self.read # type: ignore
|
||||
self.progress = Progress(
|
||||
TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
|
||||
BarColumn(bar_width=None),
|
||||
|
@ -102,9 +112,17 @@ class ProgressMonitor:
|
|||
total=self.size)
|
||||
self.progress.start()
|
||||
|
||||
def __call__(self, monitor: MultipartEncoderMonitor):
|
||||
progress = monitor.bytes_read
|
||||
self.progress.update(self.task, completed=progress)
|
||||
def read(self, size: Optional[int] = None, /) -> bytes:
|
||||
data = self._file_read(size)
|
||||
self.progress.update(self.task, completed=self.file.tell())
|
||||
return data
|
||||
|
||||
|
||||
def get_file_size(fd: BufferedIOBase) -> int:
|
||||
fd.seek(0, os.SEEK_END)
|
||||
size: int = fd.tell()
|
||||
fd.seek(0)
|
||||
return size
|
||||
|
||||
|
||||
def upload(file: Path,
|
||||
|
@ -125,30 +143,24 @@ def upload(file: Path,
|
|||
help="Check file hash before uploading."
|
||||
)):
|
||||
|
||||
if copy:
|
||||
pyperclip.copy("")
|
||||
with get_client():
|
||||
if copy:
|
||||
pyperclip.copy("")
|
||||
|
||||
if obstruct:
|
||||
filename = f"{hashlib.sha256(file.name.encode()).hexdigest()}{file.suffix}"
|
||||
elif filename is None:
|
||||
filename = file.name
|
||||
if obstruct:
|
||||
filename = f"{hashlib.sha256(file.name.encode()).hexdigest()}{file.suffix}"
|
||||
elif filename is None:
|
||||
filename = file.name
|
||||
|
||||
if hash and (resp := check_hash(filename, file)):
|
||||
return output(copy, resp)
|
||||
if hash and (resp := check_hash(filename, file)):
|
||||
return output(copy, resp)
|
||||
|
||||
with file.open('br') as f:
|
||||
files = MultipartEncoder(
|
||||
fields={'file': (filename, f)}
|
||||
)
|
||||
progress_monitor = ProgressMonitor(files.len, filename)
|
||||
with file.open('br') as f:
|
||||
monitor = ProgressMonitor(filename, f)
|
||||
files = {'file': (filename, f)}
|
||||
|
||||
monitored = MultipartEncoderMonitor(files, progress_monitor)
|
||||
headers = {'Content-Type': monitored.content_type}
|
||||
|
||||
with requests.post(HIKARI_URL,
|
||||
data=monitored, # type: ignore
|
||||
headers=headers) as req:
|
||||
progress_monitor.progress.stop()
|
||||
req = get_client().post(HIKARI_URL, files=files)
|
||||
monitor.progress.stop()
|
||||
|
||||
req.raise_for_status()
|
||||
data = UploadResponse(**req.json())
|
||||
|
|
|
@ -8,8 +8,7 @@ dynamic = ["version"]
|
|||
license = {file = "LICENSE"}
|
||||
dependencies = [
|
||||
"pyperclip",
|
||||
"requests",
|
||||
"requests-toolbelt",
|
||||
"httpx",
|
||||
"rich",
|
||||
"typer",
|
||||
]
|
||||
|
|
Loading…
Reference in a new issue