diff --git a/tox_poetry_installer/installer.py b/tox_poetry_installer/installer.py index 28fb020..6f612d6 100644 --- a/tox_poetry_installer/installer.py +++ b/tox_poetry_installer/installer.py @@ -2,7 +2,10 @@ # Silence this one globally to support the internal function imports for the proxied poetry module. # See the docstring in 'tox_poetry_installer._poetry' for more context. # pylint: disable=import-outside-toplevel +import concurrent.futures +import contextlib import typing +from typing import Optional from typing import Sequence from typing import Set @@ -18,13 +21,18 @@ if typing.TYPE_CHECKING: def install( - poetry: "_poetry.Poetry", venv: ToxVirtualEnv, packages: Sequence[PoetryPackage] + poetry: "_poetry.Poetry", + venv: ToxVirtualEnv, + packages: Sequence[PoetryPackage], + parallels: Optional[int] = None, ): """Install a bunch of packages to a virtualenv :param poetry: Poetry object the packages were sourced from :param venv: Tox virtual environment to install the packages to :param packages: List of packages to install to the virtual environment + :param parallels: Number of parallel processes to use for installing dependency packages, or + ``None`` to disable parallelization. """ from tox_poetry_installer import _poetry @@ -40,14 +48,34 @@ def install( installed: Set[PoetryPackage] = set() - for dependency in packages: - if dependency not in installed: - tox.reporter.verbosity2( - f"{constants.REPORTER_PREFIX} Installing {dependency}" - ) - pip.install(dependency) - installed.add(dependency) + @contextlib.contextmanager + def _optional_parallelize(): + """A bit of cheat, really + + A context manager that exposes a common interface for the caller that optionally + enables/disables the usage of the parallel thread pooler depending on the value of + the ``parallels`` parameter. + """ + if parallels: + with concurrent.futures.ThreadPoolExecutor( + max_workers=parallels + ) as executor: + yield executor.submit else: - tox.reporter.verbosity2( - f"{constants.REPORTER_PREFIX} Skipping {dependency}, already installed" - ) + yield lambda func, arg: func(arg) + + with _optional_parallelize() as executor: + for dependency in packages: + if dependency not in installed: + installed.add(dependency) + tox.reporter.verbosity2( + f"{constants.REPORTER_PREFIX} Installing {dependency}" + ) + executor(pip.install, dependency) + else: + tox.reporter.verbosity2( + f"{constants.REPORTER_PREFIX} Skipping {dependency}, already installed" + ) + tox.reporter.verbosity2( + f"{constants.REPORTER_PREFIX} Waiting for installs to finish..." + )