set tqdm total to None in fetch gunzip (#6377)

fixed bar overflow

python examples/beautiful_mnist.py
https://storage.googleapis.com/cvdf-datasets/mnist/train-images-idx3-ubyte.gz: 47.0MB [00:00, 136MB/s]]
https://storage.googleapis.com/cvdf-datasets/mnist/train-labels-idx1-ubyte.gz: 60.0kB [00:00, 8.66MB/s]
https://storage.googleapis.com/cvdf-datasets/mnist/t10k-images-idx3-ubyte.gz: 7.84MB [00:00, 97.5MB/s]
https://storage.googleapis.com/cvdf-datasets/mnist/t10k-labels-idx1-ubyte.gz: 10.0kB [00:00, 8.90MB/s]
This commit is contained in:
chenyu
2024-09-05 21:43:31 -04:00
committed by GitHub
parent 66e7e51c79
commit 8661276b20

View File

@@ -272,15 +272,15 @@ def fetch(url:str, name:Optional[Union[pathlib.Path, str]]=None, subdir:Optional
if not fp.is_file() or not allow_caching:
with urllib.request.urlopen(url, timeout=10) as r:
assert r.status == 200
total_length = int(r.headers.get('content-length', 0))
progress_bar = tqdm(total=total_length, unit='B', unit_scale=True, desc=f"{url}", disable=CI)
length = int(r.headers.get('content-length', 0)) if not gunzip else None
progress_bar = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI)
(path := fp.parent).mkdir(parents=True, exist_ok=True)
readfile = gzip.GzipFile(fileobj=r) if gunzip else r
with tempfile.NamedTemporaryFile(dir=path, delete=False) as f:
while chunk := readfile.read(16384): progress_bar.update(f.write(chunk))
f.close()
progress_bar.update(close=True)
if (file_size:=os.stat(f.name).st_size) < total_length: raise RuntimeError(f"fetch size incomplete, {file_size} < {total_length}")
if length and (file_size:=os.stat(f.name).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")
pathlib.Path(f.name).rename(fp)
return fp