Give batch downloader a separate class

This commit is contained in:
Nguyễn Gia Phong 2020-08-11 22:56:37 +07:00
parent a1aeb4ce01
commit b46576d933
5 changed files with 53 additions and 26 deletions

View File

@ -16,7 +16,6 @@ from pip._internal.exceptions import CommandError, PreviousBuildDirError
from pip._internal.index.collector import LinkCollector
from pip._internal.index.package_finder import PackageFinder
from pip._internal.models.selection_prefs import SelectionPreferences
from pip._internal.network.download import Downloader
from pip._internal.network.session import PipSession
from pip._internal.operations.prepare import RequirementPreparer
from pip._internal.req.constructors import (
@ -213,8 +212,6 @@ class RequirementCommand(IndexGroupCommand):
"""
Create a RequirementPreparer instance for the given parameters.
"""
downloader = Downloader(session, progress_bar=options.progress_bar)
temp_build_dir_path = temp_build_dir.path
assert temp_build_dir_path is not None
@ -239,7 +236,7 @@ class RequirementCommand(IndexGroupCommand):
build_isolation=options.build_isolation,
req_tracker=req_tracker,
session=session,
downloader=downloader,
progress_bar=options.progress_bar,
finder=finder,
require_hashes=options.require_hashes,
use_user_site=use_user_site,

View File

@ -151,7 +151,7 @@ class Downloader(object):
self._session = session
self._progress_bar = progress_bar
def download_one(self, link, location):
def __call__(self, link, location):
# type: (Link, str) -> Tuple[str, str]
"""Download the file given by link into location."""
try:
@ -173,8 +173,38 @@ class Downloader(object):
content_type = resp.headers.get('Content-Type', '')
return filepath, content_type
def download_many(self, links, location):
class BatchDownloader(object):
def __init__(
self,
session, # type: PipSession
progress_bar, # type: str
):
# type: (...) -> None
self._session = session
self._progress_bar = progress_bar
def __call__(self, links, location):
# type: (Iterable[Link], str) -> Iterable[Tuple[str, Tuple[str, str]]]
"""Download the files given by links into location."""
for link in links:
yield link.url, self.download_one(link, location)
try:
resp = _http_get_download(self._session, link)
except NetworkConnectionError as e:
assert e.response is not None
logger.critical(
"HTTP error %s while getting %s",
e.response.status_code, link,
)
raise
filename = _get_http_response_filename(resp, link)
filepath = os.path.join(location, filename)
chunks = _prepare_download(resp, link, self._progress_bar)
with open(filepath, 'wb') as content_file:
for chunk in chunks:
content_file.write(chunk)
content_type = resp.headers.get('Content-Type', '')
yield link.url, (filepath, content_type)

View File

@ -26,6 +26,7 @@ from pip._internal.exceptions import (
VcsHashUnsupported,
)
from pip._internal.models.wheel import Wheel
from pip._internal.network.download import BatchDownloader, Downloader
from pip._internal.network.lazy_wheel import (
HTTPRangeRequestUnsupported,
dist_from_wheel_url,
@ -52,7 +53,6 @@ if MYPY_CHECK_RUNNING:
from pip._internal.index.package_finder import PackageFinder
from pip._internal.models.link import Link
from pip._internal.network.download import Downloader
from pip._internal.network.session import PipSession
from pip._internal.req.req_install import InstallRequirement
from pip._internal.req.req_tracker import RequirementTracker
@ -116,7 +116,7 @@ class File(object):
def get_http_url(
link, # type: Link
downloader, # type: Downloader
download, # type: Downloader
download_dir=None, # type: Optional[str]
hashes=None, # type: Optional[Hashes]
):
@ -134,7 +134,7 @@ def get_http_url(
content_type = None
else:
# let's download to a tmp dir
from_path, content_type = downloader.download_one(link, temp_dir.path)
from_path, content_type = download(link, temp_dir.path)
if hashes:
hashes.check_against_path(from_path)
@ -227,7 +227,7 @@ def get_file_url(
def unpack_url(
link, # type: Link
location, # type: str
downloader, # type: Downloader
download, # type: Downloader
download_dir=None, # type: Optional[str]
hashes=None, # type: Optional[Hashes]
):
@ -259,7 +259,7 @@ def unpack_url(
else:
file = get_http_url(
link,
downloader,
download,
download_dir,
hashes=hashes,
)
@ -311,7 +311,7 @@ class RequirementPreparer(object):
build_isolation, # type: bool
req_tracker, # type: RequirementTracker
session, # type: PipSession
downloader, # type: Downloader
progress_bar, # type: str
finder, # type: PackageFinder
require_hashes, # type: bool
use_user_site, # type: bool
@ -324,7 +324,8 @@ class RequirementPreparer(object):
self.build_dir = build_dir
self.req_tracker = req_tracker
self._session = session
self.downloader = downloader
self._download = Downloader(session, progress_bar)
self._batch_download = BatchDownloader(session, progress_bar)
self.finder = finder
# Where still-packed archives should be written to. If None, they are
@ -508,7 +509,7 @@ class RequirementPreparer(object):
# Let's download to a temporary directory.
tmpdir = TempDirectory(kind="unpack", globally_managed=True).path
self._downloaded.update(self.downloader.download_many(links, tmpdir))
self._downloaded.update(self._batch_download(links, tmpdir))
for req in reqs:
self._prepare_linked_requirement(req, parallel_builds)
@ -524,7 +525,7 @@ class RequirementPreparer(object):
if link.url not in self._downloaded:
try:
local_file = unpack_url(
link, req.source_dir, self.downloader,
link, req.source_dir, self._download,
download_dir, hashes,
)
except NetworkConnectionError as exc:

View File

@ -35,7 +35,7 @@ def test_unpack_url_with_urllib_response_without_content_type(data):
session = Mock()
session.get = _fake_session_get
downloader = Downloader(session, progress_bar="on")
download = Downloader(session, progress_bar="on")
uri = path_to_url(data.packages.joinpath("simple-1.0.tar.gz"))
link = Link(uri)
@ -44,7 +44,7 @@ def test_unpack_url_with_urllib_response_without_content_type(data):
unpack_url(
link,
temp_dir,
downloader=downloader,
download=download,
download_dir=None,
)
assert set(os.listdir(temp_dir)) == {
@ -75,11 +75,11 @@ def test_download_http_url__no_directory_traversal(mock_raise_for_status,
'content-disposition': 'attachment;filename="../out_dir_file"'
}
session.get.return_value = resp
downloader = Downloader(session, progress_bar="on")
download = Downloader(session, progress_bar="on")
download_dir = tmpdir.joinpath('download')
os.mkdir(download_dir)
file_path, content_type = downloader.download_one(link, download_dir)
file_path, content_type = download(link, download_dir)
# The file should be downloaded to download_dir.
actual = os.listdir(download_dir)
assert actual == ['out_dir_file']
@ -178,11 +178,11 @@ class Test_unpack_url(object):
self.dist_path2 = data.packages.joinpath(self.dist_file2)
self.dist_url = Link(path_to_url(self.dist_path))
self.dist_url2 = Link(path_to_url(self.dist_path2))
self.no_downloader = Mock(side_effect=AssertionError)
self.no_download = Mock(side_effect=AssertionError)
def test_unpack_url_no_download(self, tmpdir, data):
self.prep(tmpdir, data)
unpack_url(self.dist_url, self.build_dir, self.no_downloader)
unpack_url(self.dist_url, self.build_dir, self.no_download)
assert os.path.isdir(os.path.join(self.build_dir, 'simple'))
assert not os.path.isfile(
os.path.join(self.download_dir, self.dist_file))
@ -198,7 +198,7 @@ class Test_unpack_url(object):
with pytest.raises(HashMismatch):
unpack_url(dist_url,
self.build_dir,
downloader=self.no_downloader,
download=self.no_download,
hashes=Hashes({'md5': ['bogus']}))
def test_unpack_url_thats_a_dir(self, tmpdir, data):
@ -206,7 +206,7 @@ class Test_unpack_url(object):
dist_path = data.packages.joinpath("FSPkg")
dist_url = Link(path_to_url(dist_path))
unpack_url(dist_url, self.build_dir,
downloader=self.no_downloader,
download=self.no_download,
download_dir=self.download_dir)
assert os.path.isdir(os.path.join(self.build_dir, 'fspkg'))

View File

@ -18,7 +18,6 @@ from pip._internal.exceptions import (
InvalidWheelFilename,
PreviousBuildDirError,
)
from pip._internal.network.download import Downloader
from pip._internal.network.session import PipSession
from pip._internal.operations.prepare import RequirementPreparer
from pip._internal.req import InstallRequirement, RequirementSet
@ -87,7 +86,7 @@ class TestRequirementSet(object):
build_isolation=True,
req_tracker=tracker,
session=session,
downloader=Downloader(session, progress_bar="on"),
progress_bar='on',
finder=finder,
require_hashes=require_hashes,
use_user_site=False,