优化 util 逻辑

BiliDriveEx/__main__.py

@@ -143,7 +143,7 @@ def upload_handle(args):
         log("暂不支持上传文件夹")
         return None
     log(f"上传: {os.path.basename(file_name)} ({size_string(os.path.getsize(file_name))})")
-    first_4mb_sha1 = calc_sha1(read_in_chunk(file_name, chunk_size=4 * 1024 * 1024, chunk_number=1))
+    first_4mb_sha1 = calc_sha1(read_in_chunk(file_name, size=4 * 1024 * 1024, cnt=1))
     history = read_history()
     if first_4mb_sha1 in history:
         url = history[first_4mb_sha1]['url']
@@ -162,7 +162,7 @@ def upload_handle(args):
     thread_pool = []
     block_dict = {}
     block_num = math.ceil(os.path.getsize(file_name) / (args.block_size * 1024 * 1024))
-    for index, block in enumerate(read_in_chunk(file_name, chunk_size=args.block_size * 1024 * 1024)):
+    for index, block in enumerate(read_in_chunk(file_name, size=args.block_size * 1024 * 1024)):
         if len(thread_pool) >= args.thread:
             done_flag.acquire()
         if not terminate_flag.is_set():

BiliDriveEx/util.py

@@ -52,16 +52,15 @@ def read_history():
     return history
     
 
-def read_in_chunk(file_name, chunk_size=4 * 1024 * 1024, chunk_number=-1):
-    chunk_counter = 0
-    with open(file_name, "rb") as f:
+def read_in_chunk(fname, size=4 * 1024 * 1024, cnt=-1):
+    with open(fname, "rb") as f:
+        idx = 0
         while True:
-            data = f.read(chunk_size)
-            if data != b"" and (chunk_number == -1 or chunk_counter < chunk_number):
-                yield data
-                chunk_counter += 1
-            else:
-                return
+            data = f.read(size)
+            if not data or (cnt != -1 and idx >= cnt):
+                break
+            yield data
+            idx += 1
                 
 def log(message):
     print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}")
\ No newline at end of file