Skip to content

Commit

Permalink
cleanup download utils (#8273)
Browse files Browse the repository at this point in the history
  • Loading branch information
pmeier authored Feb 15, 2024
1 parent f21c6bd commit b1123cf
Showing 1 changed file with 5 additions and 17 deletions.
22 changes: 5 additions & 17 deletions torchvision/datasets/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import urllib.error
import urllib.request
import zipfile
from typing import Any, Callable, Dict, IO, Iterable, Iterator, List, Optional, Tuple, TypeVar, Union
from typing import Any, Callable, Dict, IO, Iterable, List, Optional, Tuple, TypeVar, Union
from urllib.parse import urlparse

import numpy as np
Expand All @@ -24,24 +24,12 @@
USER_AGENT = "pytorch/vision"


def _save_response_content(
content: Iterator[bytes],
destination: Union[str, pathlib.Path],
length: Optional[int] = None,
) -> None:
with open(destination, "wb") as fh, tqdm(total=length) as pbar:
for chunk in content:
# filter out keep-alive new chunks
if not chunk:
continue

fh.write(chunk)
pbar.update(len(chunk))


def _urlretrieve(url: str, filename: Union[str, pathlib.Path], chunk_size: int = 1024 * 32) -> None:
with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT})) as response:
_save_response_content(iter(lambda: response.read(chunk_size), b""), filename, length=response.length)
with open(filename, "wb") as fh, tqdm(total=response.length) as pbar:
while chunk := response.read(chunk_size):
fh.write(chunk)
pbar.update(len(chunk))


def calculate_md5(fpath: Union[str, pathlib.Path], chunk_size: int = 1024 * 1024) -> str:
Expand Down

0 comments on commit b1123cf

Please sign in to comment.