transbeam-cli: load files asynchronously

This commit is contained in:
xenofem 2024-04-03 13:01:51 -04:00
parent 60d9636896
commit 35dca63cff
3 changed files with 17 additions and 7 deletions

View file

@ -3,9 +3,10 @@
from setuptools import setup from setuptools import setup
setup(name='transbeam-cli', setup(name='transbeam-cli',
version='0.1', version='0.2',
scripts=['transbeam-cli'], scripts=['transbeam-cli'],
install_requires=[ install_requires=[
'aiofiles',
'tqdm', 'tqdm',
'websockets', 'websockets',
], ],

View file

@ -7,17 +7,24 @@ import json
import math import math
import pathlib import pathlib
import sys import sys
import aiofiles
from tqdm import tqdm from tqdm import tqdm
from websockets import connect from websockets import connect
def file_loader(files): FILE_CHUNK_SIZE = 16384
async def file_loader(files):
with tqdm(desc="Total", total=sum(size for (path, size) in files), unit='B', unit_scale=True, leave=True, position=1) as total_progress: with tqdm(desc="Total", total=sum(size for (path, size) in files), unit='B', unit_scale=True, leave=True, position=1) as total_progress:
for (path, size) in files: for (path, size) in files:
with tqdm(desc=path.name, total=size, unit='B', unit_scale=True, leave=True, position=0) as file_progress: with tqdm(desc=path.name, total=size, unit='B', unit_scale=True, leave=True, position=0) as file_progress:
with path.open("rb") as f: async with aiofiles.open(path, mode='rb') as f:
while f.tell() < size: while True:
data = f.read(min(16384, size - f.tell())) pos = await f.tell()
if data == "": if pos >= size:
break
data = await f.read(min(FILE_CHUNK_SIZE, size - pos))
if data == b'':
tqdm.write("file ended early!") tqdm.write("file ended early!")
exit(1) exit(1)
total_progress.update(len(data)) total_progress.update(len(data))
@ -49,7 +56,7 @@ async def send(paths, host, password, lifetime, collection_name=None):
exit(1) exit(1)
print("Download: https://{}/download?code={}".format(host, resp["code"])) print("Download: https://{}/download?code={}".format(host, resp["code"]))
loader = file_loader([(paths[i], fileMetadata[i]["size"]) for i in range(len(paths))]) loader = file_loader([(paths[i], fileMetadata[i]["size"]) for i in range(len(paths))])
for data in loader: async for data in loader:
await ws.send(data) await ws.send(data)
parser = argparse.ArgumentParser(description="Upload files to transbeam") parser = argparse.ArgumentParser(description="Upload files to transbeam")

View file

@ -49,6 +49,7 @@
pname = "${name}-cli"; pname = "${name}-cli";
version = "0.1"; version = "0.1";
propagatedBuildInputs = [ propagatedBuildInputs = [
aiofiles
tqdm tqdm
websockets websockets
]; ];
@ -70,6 +71,7 @@
cargo-audit cargo-audit
cargo-flamegraph cargo-flamegraph
(python3.withPackages (p: with p; [ (python3.withPackages (p: with p; [
aiofiles
tqdm tqdm
websockets websockets
])) ]))