diff --git a/hb/main.py b/hb/main.py index 3c6810f..070d6e7 100644 --- a/hb/main.py +++ b/hb/main.py @@ -7,7 +7,7 @@ from sys import stderr from threading import Thread from time import sleep -from typing import Dict, IO, List, Tuple +from typing import Dict, IO, List, Optional, Tuple class Checksum(): @@ -44,21 +44,26 @@ def __init__(self, path: str, threshold: int = 200) -> None: self.filesize = self._path.stat().st_size self.threshold = threshold - def _progress(self, file: IO) -> None: + def _progress(self, file: IO) -> Optional[Thread]: def _p(file: IO) -> None: while not file.closed: print(f"{int(file.tell() / self.filesize * 100)}%", end="\r", file=stderr) sleep(0.2) print(" ", end="\r", file=stderr) # clear the progress display + thread = None if self.filesize > self.threshold * 1024 * 1024: - Thread(target=_p, args=(file,)).start() + thread = Thread(target=_p, args=(file,)) + thread.start() + return thread def _hashlib_compute(self, name: str) -> str: result = hashlib.new(name) with self._path.open("rb") as lines: - self._progress(lines) + thread = self._progress(lines) for line in lines: result.update(line) + if thread: + thread.join() return result.hexdigest() def _zlib_compute(self, name: str) -> str: @@ -69,9 +74,11 @@ def _zlib_compute(self, name: str) -> str: result = 0 update = zlib.crc32 with self._path.open("rb") as lines: - self._progress(lines) + thread = self._progress(lines) for line in lines: result = update(line, result) + if thread: + thread.join() return hex(result)[2:].zfill(8) def compute(self, algorithm: str) -> str: