优化统计

fast_media_lock.py

@@ -214,11 +214,15 @@ class _ProgressWriter:
         self._enabled = sys.stdout.isatty()
         self._lock = threading.Lock()
         self._last_len = 0
+        self._file_mode = False
 
     def update(self, label: str, done: int, total: int, path: Path) -> None:
         if not self._enabled:
             return
 
+        if not self._file_mode:
+            return
+
         total = max(total, 1)
         percent = (done * 100) // total
         width = 24
@@ -236,15 +240,48 @@ class _ProgressWriter:
         if not self._enabled:
             return
 
+        if not self._file_mode:
+            return
+
         with self._lock:
             sys.stdout.write("\n")
             sys.stdout.flush()
             self._last_len = 0
 
+    def finish_count(self) -> None:
+        if not self._enabled:
+            return
+
+        with self._lock:
+            sys.stdout.write("\n")
+            sys.stdout.flush()
+            self._last_len = 0
+
+    def set_file_mode(self, enabled: bool) -> None:
+        self._file_mode = enabled
+
+    def count_update(self, done: int, total: int) -> None:
+        if not self._enabled:
+            return
+        with self._lock:
+            msg = f"已完成: {done}/{total}"
+            pad = max(0, self._last_len - len(msg))
+            sys.stdout.write("\r" + msg + (" " * pad))
+            sys.stdout.flush()
+            self._last_len = len(msg)
+
 
 _progress = _ProgressWriter()
 
 
+def _iter_files_list(target: Path, all_files: bool) -> list[Path]:
+    return list(_iter_files(target, all_files=all_files))
+
+
+def _iter_files_by_names_list(target: Path, names: list[str], all_files: bool) -> list[Path]:
+    return list(_iter_files_by_names(target, names, all_files=all_files))
+
+
 def encrypt_file(path: Path, password: str, chunk_size: int) -> str:
     if not path.exists() or not path.is_file():
         return "skip(not_file)"
@@ -370,63 +407,145 @@ def _resolve_workers(workers: int) -> int:
     return workers
 
 
-def run_encrypt(target: Path, password: str, chunk_mb: int, all_files: bool, workers: int, names: list[str] | None = None) -> int:
+def run_encrypt(target: Path, password: str, chunk_mb: int, all_files: bool, workers: int, file_progress: bool, names: list[str] | None = None) -> int:
     chunk_size = chunk_mb * 1024 * 1024
     worker_count = _resolve_workers(workers)
     ok = 0
     skipped = 0
     failed = 0
-
-    if worker_count == 1:
-        results = ((p, encrypt_file(p, password, chunk_size)) for p in (_iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)))
+    _progress.set_file_mode(file_progress)
+
+    files = _iter_files_by_names_list(target, names, all_files=all_files) if names else _iter_files_list(target, all_files=all_files)
+    total_files = len(files)
+
+    if not file_progress:
+        if worker_count == 1:
+            completed = 0
+            results = []
+            try:
+                for p in files:
+                    results.append((p, encrypt_file(p, password, chunk_size)))
+                    completed += 1
+                    _progress.count_update(completed, total_files)
+            except KeyboardInterrupt:
+                print("\n已收到中断信号，当前文件处理完成后退出", file=sys.stderr)
+        else:
+            completed = 0
+            with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
+                future_map = {ex.submit(encrypt_file, p, password, chunk_size): p for p in files}
+                results = []
+                try:
+                    for fut in concurrent.futures.as_completed(future_map):
+                        p = future_map[fut]
+                        res = fut.result()
+                        results.append((p, res))
+                        completed += 1
+                        _progress.count_update(completed, total_files)
+                except KeyboardInterrupt:
+                    print("\n已收到中断信号，等待当前文件处理完成后退出", file=sys.stderr)
+                    for fut, p in future_map.items():
+                        if fut.done():
+                            continue
+                        try:
+                            res = fut.result()
+                        except Exception as e:
+                            res = f"fail({type(e).__name__})"
+                        results.append((p, res))
     else:
-        with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
-            file_iter = _iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)
-            results = ex.map(_encrypt_task, ((p, password, chunk_size) for p in file_iter))
+        if worker_count == 1:
+            results = ((p, encrypt_file(p, password, chunk_size)) for p in files)
+        else:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
+                results = ex.map(_encrypt_task, ((p, password, chunk_size) for p in files))
 
     for p, res in results:
         if res == "ok":
             ok += 1
-            if not sys.stdout.isatty():
+            if not file_progress and not sys.stdout.isatty():
                 print(f"[OK] encrypted: {p}")
         elif res.startswith("fail"):
             failed += 1
-            print(f"[FAIL] {p} -> {res}")
+            if not file_progress:
+                print(f"[FAIL] {p} -> {res}")
         else:
             skipped += 1
-            if not sys.stdout.isatty():
+            if not file_progress and not sys.stdout.isatty():
                 print(f"[SKIP] {p} -> {res}")
 
+    if not file_progress:
+        _progress.finish_count()
+
     print(f"完成: encrypted={ok}, failed={failed}, skipped={skipped}, chunk={chunk_mb}MB, workers={worker_count}")
     return 2 if failed else 0
 
 
-def run_decrypt(target: Path, password: str, all_files: bool, workers: int, names: list[str] | None = None) -> int:
+def run_decrypt(target: Path, password: str, all_files: bool, workers: int, file_progress: bool, names: list[str] | None = None) -> int:
     worker_count = _resolve_workers(workers)
     ok = 0
     skipped = 0
     failed = 0
-
-    if worker_count == 1:
-        results = ((p, decrypt_file(p, password)) for p in (_iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)))
+    _progress.set_file_mode(file_progress)
+
+    files = _iter_files_by_names_list(target, names, all_files=all_files) if names else _iter_files_list(target, all_files=all_files)
+    total_files = len(files)
+
+    if not file_progress:
+        if worker_count == 1:
+            completed = 0
+            results = []
+            try:
+                for p in files:
+                    results.append((p, decrypt_file(p, password)))
+                    completed += 1
+                    _progress.count_update(completed, total_files)
+            except KeyboardInterrupt:
+                print("\n已收到中断信号，当前文件处理完成后退出", file=sys.stderr)
+        else:
+            completed = 0
+            with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
+                future_map = {ex.submit(decrypt_file, p, password): p for p in files}
+                results = []
+                try:
+                    for fut in concurrent.futures.as_completed(future_map):
+                        p = future_map[fut]
+                        res = fut.result()
+                        results.append((p, res))
+                        completed += 1
+                        _progress.count_update(completed, total_files)
+                except KeyboardInterrupt:
+                    print("\n已收到中断信号，等待当前文件处理完成后退出", file=sys.stderr)
+                    for fut, p in future_map.items():
+                        if fut.done():
+                            continue
+                        try:
+                            res = fut.result()
+                        except Exception as e:
+                            res = f"fail({type(e).__name__})"
+                        results.append((p, res))
     else:
-        with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
-            file_iter = _iter_files_by_names(target, names, all_files=all_files) if names else _iter_files(target, all_files=all_files)
-            results = ex.map(_decrypt_task, ((p, password) for p in file_iter))
+        if worker_count == 1:
+            results = ((p, decrypt_file(p, password)) for p in files)
+        else:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as ex:
+                results = ex.map(_decrypt_task, ((p, password) for p in files))
 
     for p, res in results:
         if res == "ok":
             ok += 1
-            if not sys.stdout.isatty():
+            if not file_progress and not sys.stdout.isatty():
                 print(f"[OK] decrypted: {p}")
         elif res.startswith("fail"):
             failed += 1
-            print(f"[FAIL] {p} -> {res}")
+            if not file_progress:
+                print(f"[FAIL] {p} -> {res}")
         else:
             skipped += 1
-            if not sys.stdout.isatty():
+            if not file_progress and not sys.stdout.isatty():
                 print(f"[SKIP] {p} -> {res}")
 
+    if not file_progress:
+        _progress.finish_count()
+
     print(f"完成: decrypted={ok}, failed={failed}, skipped={skipped}, workers={worker_count}")
     return 2 if failed else 0
 
@@ -467,6 +586,7 @@ def build_parser() -> argparse.ArgumentParser:
         if name in ("lock-name", "unlock-name"):
             p.add_argument("names", nargs="+", help="文件名或相对路径（可一次传多个）")
         p.add_argument("--media-only", action="store_true", help="仅处理媒体后缀（默认处理所有文件）")
+        p.add_argument("--file-progress", action="store_true", help="显示单文件进度条，默认显示已完成文件数/总文件数")
         if name in ("lock", "lock-name"):
             p.add_argument("--chunk-mb", type=int, default=1, help="加密前多少MB（默认8）")
         if name in ("lock", "unlock", "lock-name", "unlock-name"):
@@ -492,21 +612,21 @@ def main() -> int:
             if args.chunk_mb <= 0:
                 raise LockerError("--chunk-mb 必须大于0")
             pw = args.password if args.password else _ask_password(confirm=True)
-            return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers)
+            return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers, file_progress=args.file_progress)
 
         if args.cmd in ("unlock", "decrypt"):
             pw = args.password if args.password else _ask_password(confirm=False)
-            return run_decrypt(target, pw, all_files=all_files, workers=args.workers)
+            return run_decrypt(target, pw, all_files=all_files, workers=args.workers, file_progress=args.file_progress)
 
         if args.cmd in ("lock-name", "encrypt-name"):
             if args.chunk_mb <= 0:
                 raise LockerError("--chunk-mb 必须大于0")
             pw = args.password if args.password else _ask_password(confirm=True)
-            return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers, names=args.names)
+            return run_encrypt(target, pw, args.chunk_mb, all_files=all_files, workers=args.workers, file_progress=args.file_progress, names=args.names)
 
         if args.cmd in ("unlock-name", "decrypt-name"):
             pw = args.password if args.password else _ask_password(confirm=False)
-            return run_decrypt(target, pw, all_files=all_files, workers=args.workers, names=args.names)
+            return run_decrypt(target, pw, all_files=all_files, workers=args.workers, file_progress=args.file_progress, names=args.names)
 
         return run_status(target, all_files=all_files)