use httpx instead of requests

httpx is already typed, upload monitoring isn't so great
This commit is contained in:
odrling 2024-03-22 01:16:57 +01:00
parent a79585f52f
commit 2254ee3ec4
Signed by: odrling
GPG key ID: 2D8C2F151EDB7392
2 changed files with 46 additions and 35 deletions

View file

@ -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())

View file

@ -8,8 +8,7 @@ dynamic = ["version"]
license = {file = "LICENSE"}
dependencies = [
"pyperclip",
"requests",
"requests-toolbelt",
"httpx",
"rich",
"typer",
]