diff --git a/pex/commands/command.py b/pex/commands/command.py index b32330237..32575acb6 100644 --- a/pex/commands/command.py +++ b/pex/commands/command.py @@ -17,7 +17,7 @@ from pex import pex_warnings from pex.argparse import HandleBoolAction from pex.cache import access as cache_access -from pex.common import safe_mkdtemp, safe_open +from pex.common import environment_as, safe_mkdtemp, safe_open from pex.compatibility import shlex_quote from pex.result import Error, Ok, Result from pex.typing import TYPE_CHECKING, Generic, cast @@ -379,7 +379,7 @@ def warn_ignore_pex_root(set_via): else: pex_root = options.cache_dir or options.pex_root or ENV.PEX_ROOT - with ENV.patch(PEX_ROOT=pex_root, TMPDIR=tmpdir) as env: + with ENV.patch(PEX_ROOT=pex_root, TMPDIR=tmpdir) as env, environment_as(**env): cache_access.read_write() yield env diff --git a/pex/common.py b/pex/common.py index 1135dee9e..056ce6fbf 100644 --- a/pex/common.py +++ b/pex/common.py @@ -17,6 +17,7 @@ import time import zipfile from collections import defaultdict, namedtuple +from contextlib import contextmanager from datetime import datetime from uuid import uuid4 from zipfile import ZipFile, ZipInfo @@ -917,3 +918,28 @@ def iter_copytree( if copy_mode is CopyMode.SYMLINK: # Once we've symlinked the top-level directories and files, we've "copied" everything. return + + +@contextmanager +def environment_as(**kwargs): + # type: (**Any) -> Iterator[None] + """Mutates the `os.environ` for the duration of the context. + + Keyword arguments with None values are removed from os.environ (if present) and all other + keyword arguments are added or updated in `os.environ` with the values taken from the + stringification (`str(...)`) of each value. + """ + existing = {key: os.environ.get(key) for key in kwargs} + + def adjust_environment(mapping): + for key, value in mapping.items(): + if value is not None: + os.environ[key] = str(value) + else: + os.environ.pop(key, None) + + adjust_environment(kwargs) + try: + yield + finally: + adjust_environment(existing) diff --git a/pex/jobs.py b/pex/jobs.py index 2fa1234f4..67513596b 100644 --- a/pex/jobs.py +++ b/pex/jobs.py @@ -17,7 +17,7 @@ from pex.common import pluralize from pex.compatibility import Queue, cpu_count from pex.tracer import TRACER -from pex.typing import TYPE_CHECKING, Generic, cast +from pex.typing import TYPE_CHECKING, Generic if TYPE_CHECKING: from typing import ( @@ -679,12 +679,8 @@ def join(self): @contextmanager def _mp_pool(size): # type: (int) -> Iterator[Pool] - try: - context = multiprocessing.get_context("fork") # type: ignore[attr-defined] - pool = cast("Pool", context.Pool(processes=size)) - except (AttributeError, ValueError): - pool = multiprocessing.Pool(processes=size) + pool = multiprocessing.Pool(processes=size) try: yield pool finally: diff --git a/testing/__init__.py b/testing/__init__.py index 4e22f4d8a..b9be3f12a 100644 --- a/testing/__init__.py +++ b/testing/__init__.py @@ -674,25 +674,6 @@ def all_python_venvs(system_site_packages=False): ) -@contextmanager -def environment_as(**kwargs): - # type: (**Any) -> Iterator[None] - existing = {key: os.environ.get(key) for key in kwargs} - - def adjust_environment(mapping): - for key, value in mapping.items(): - if value is not None: - os.environ[key] = str(value) - else: - os.environ.pop(key, None) - - adjust_environment(kwargs) - try: - yield - finally: - adjust_environment(existing) - - @contextmanager def pushd(directory): # type: (Text) -> Iterator[None] diff --git a/tests/integration/cli/commands/test_cache_prune.py b/tests/integration/cli/commands/test_cache_prune.py index 4049c012a..788bd9624 100644 --- a/tests/integration/cli/commands/test_cache_prune.py +++ b/tests/integration/cli/commands/test_cache_prune.py @@ -27,13 +27,13 @@ VenvDirs, ) from pex.cli.commands.cache.du import DiskUsage -from pex.common import safe_open +from pex.common import environment_as, safe_open from pex.pep_503 import ProjectName from pex.pex_info import PexInfo from pex.pip.version import PipVersion, PipVersionValue from pex.typing import TYPE_CHECKING from pex.variables import ENV -from testing import environment_as, run_pex_command +from testing import run_pex_command from testing.cli import run_pex3 from testing.pytest.tmp import Tempdir diff --git a/tests/integration/test_integration.py b/tests/integration/test_integration.py index bea301fe3..164763655 100644 --- a/tests/integration/test_integration.py +++ b/tests/integration/test_integration.py @@ -19,7 +19,15 @@ from pex import targets from pex.cache.dirs import CacheDir, InterpreterDir -from pex.common import is_exe, safe_mkdir, safe_open, safe_rmtree, temporary_dir, touch +from pex.common import ( + environment_as, + is_exe, + safe_mkdir, + safe_open, + safe_rmtree, + temporary_dir, + touch, +) from pex.compatibility import WINDOWS, commonpath from pex.dist_metadata import Distribution, Requirement, is_wheel from pex.fetcher import URLFetcher @@ -46,7 +54,6 @@ IntegResults, built_wheel, ensure_python_interpreter, - environment_as, get_dep_dist_names_from_pex, make_env, run_pex_command, diff --git a/tests/integration/test_issue_157.py b/tests/integration/test_issue_157.py index f9548ace9..9e17a677d 100644 --- a/tests/integration/test_issue_157.py +++ b/tests/integration/test_issue_157.py @@ -14,10 +14,11 @@ import pytest from colors import color # vendor:skip +from pex.common import environment_as from pex.pex_info import PexInfo from pex.typing import TYPE_CHECKING from pex.version import __version__ -from testing import IS_PYPY, environment_as, make_env, run_pex_command, scie +from testing import IS_PYPY, make_env, run_pex_command, scie if TYPE_CHECKING: from typing import Any, Iterable, Iterator, List, Tuple diff --git a/tests/resolve/test_target_options.py b/tests/resolve/test_target_options.py index 832cbb7ae..2a75bacbf 100644 --- a/tests/resolve/test_target_options.py +++ b/tests/resolve/test_target_options.py @@ -10,6 +10,7 @@ import pytest import pex.resolve.target_configuration +from pex.common import environment_as from pex.interpreter import PythonInterpreter from pex.pep_425 import CompatibilityTags from pex.pep_508 import MarkerEnvironment @@ -20,7 +21,7 @@ from pex.targets import CompletePlatform, Targets from pex.typing import TYPE_CHECKING from pex.variables import ENV -from testing import IS_MAC, environment_as +from testing import IS_MAC if TYPE_CHECKING: from typing import Any, Dict, Iterable, List, Optional, Tuple, Type diff --git a/tests/test_atomic_directory.py b/tests/test_atomic_directory.py index 771b37ceb..c2060b264 100644 --- a/tests/test_atomic_directory.py +++ b/tests/test_atomic_directory.py @@ -9,9 +9,8 @@ import pytest from pex.atomic_directory import AtomicDirectory, FileLockStyle, _is_bsd_lock, atomic_directory -from pex.common import temporary_dir, touch +from pex.common import environment_as, temporary_dir, touch from pex.typing import TYPE_CHECKING -from testing import environment_as try: from unittest import mock diff --git a/tests/test_interpreter.py b/tests/test_interpreter.py index 8a8a80bda..de95e0661 100644 --- a/tests/test_interpreter.py +++ b/tests/test_interpreter.py @@ -14,7 +14,7 @@ import pytest from pex.cache.dirs import InterpreterDir -from pex.common import chmod_plus_x, safe_mkdir, safe_mkdtemp, temporary_dir, touch +from pex.common import chmod_plus_x, environment_as, safe_mkdir, safe_mkdtemp, temporary_dir, touch from pex.executor import Executor from pex.interpreter import PythonInterpreter, create_shebang from pex.jobs import Job @@ -30,7 +30,6 @@ ensure_python_distribution, ensure_python_interpreter, ensure_python_venv, - environment_as, pushd, ) from testing.pytest.tmp import TempdirFactory diff --git a/tests/test_pex.py b/tests/test_pex.py index 81a2d0b0c..3a0a55ac6 100644 --- a/tests/test_pex.py +++ b/tests/test_pex.py @@ -16,7 +16,7 @@ import pytest from pex import resolver -from pex.common import safe_mkdir, safe_open, temporary_dir +from pex.common import environment_as, safe_mkdir, safe_open, temporary_dir from pex.compatibility import PY2, WINDOWS, to_bytes from pex.dist_metadata import Distribution from pex.interpreter import PythonIdentity, PythonInterpreter @@ -33,7 +33,6 @@ PY_VER, WheelBuilder, ensure_python_interpreter, - environment_as, install_wheel, make_bdist, run_simple_pex, diff --git a/tests/test_pip.py b/tests/test_pip.py index ca66c9a3f..172ab2378 100644 --- a/tests/test_pip.py +++ b/tests/test_pip.py @@ -13,7 +13,7 @@ import pytest -from pex.common import safe_rmtree +from pex.common import environment_as, safe_rmtree from pex.dist_metadata import Distribution, Requirement from pex.interpreter import PythonInterpreter from pex.jobs import Job @@ -29,7 +29,7 @@ from pex.typing import TYPE_CHECKING from pex.variables import ENV from pex.venv.virtualenv import Virtualenv -from testing import IS_LINUX, PY310, ensure_python_interpreter, environment_as +from testing import IS_LINUX, PY310, ensure_python_interpreter from testing.pytest.tmp import Tempdir if TYPE_CHECKING: diff --git a/tests/test_requirements.py b/tests/test_requirements.py index db9bb3348..7a53f0c6c 100644 --- a/tests/test_requirements.py +++ b/tests/test_requirements.py @@ -8,7 +8,7 @@ import pytest -from pex.common import safe_open, temporary_dir, touch +from pex.common import environment_as, safe_open, temporary_dir, touch from pex.dist_metadata import Requirement from pex.fetcher import URLFetcher from pex.requirements import ( @@ -30,7 +30,6 @@ ) from pex.third_party.packaging.markers import Marker from pex.typing import TYPE_CHECKING -from testing import environment_as if TYPE_CHECKING: from typing import Any, Iterable, Iterator, List, Optional, Union diff --git a/tests/test_variables.py b/tests/test_variables.py index 88e9c2277..25b85e21a 100644 --- a/tests/test_variables.py +++ b/tests/test_variables.py @@ -7,12 +7,12 @@ import pytest from pex import pex_warnings +from pex.common import environment_as from pex.compatibility import PY2 from pex.pex_warnings import PEXWarning from pex.typing import TYPE_CHECKING from pex.util import named_temporary_file from pex.variables import NoValueError, Variables -from testing import environment_as from testing.pytest.tmp import Tempdir if TYPE_CHECKING: