transbeam-cli: load files asynchronously

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',
       ],

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")

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
               ]))