diff --git a/cli/setup.py b/cli/setup.py index fe50274..911c6f4 100644 --- a/cli/setup.py +++ b/cli/setup.py @@ -3,9 +3,10 @@ from setuptools import setup setup(name='transbeam-cli', - version='0.1', + version='0.2', scripts=['transbeam-cli'], install_requires=[ + 'aiofiles', 'tqdm', 'websockets', ], diff --git a/cli/transbeam-cli b/cli/transbeam-cli index 331a725..8ee4694 100755 --- a/cli/transbeam-cli +++ b/cli/transbeam-cli @@ -7,17 +7,24 @@ import json import math import pathlib import sys + +import aiofiles from tqdm import tqdm 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: 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 path.open("rb") as f: - while f.tell() < size: - data = f.read(min(16384, size - f.tell())) - if data == "": + async with aiofiles.open(path, mode='rb') as f: + while True: + pos = await f.tell() + if pos >= size: + break + data = await f.read(min(FILE_CHUNK_SIZE, size - pos)) + if data == b'': tqdm.write("file ended early!") exit(1) total_progress.update(len(data)) @@ -49,7 +56,7 @@ async def send(paths, host, password, lifetime, collection_name=None): exit(1) print("Download: https://{}/download?code={}".format(host, resp["code"])) 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) parser = argparse.ArgumentParser(description="Upload files to transbeam") diff --git a/flake.nix b/flake.nix index 8cf65b0..0dc5b27 100644 --- a/flake.nix +++ b/flake.nix @@ -49,6 +49,7 @@ pname = "${name}-cli"; version = "0.1"; propagatedBuildInputs = [ + aiofiles tqdm websockets ]; @@ -70,6 +71,7 @@ cargo-audit cargo-flamegraph (python3.withPackages (p: with p; [ + aiofiles tqdm websockets ]))