|
| 1 | +import contextlib |
1 | 2 | import os
|
2 | 3 | import sys
|
3 | 4 | from typing import Union
|
4 | 5 |
|
| 6 | +from more_itertools import unique_everseen |
| 7 | + |
| 8 | + |
5 | 9 | if sys.version_info >= (3, 9):
|
6 | 10 | StrPath = Union[str, os.PathLike[str]] # Same as _typeshed.StrPath
|
7 | 11 | else:
|
@@ -38,3 +42,41 @@ def normpath(filename: StrPath) -> str:
|
38 | 42 | # See pkg_resources.normalize_path for notes about cygwin
|
39 | 43 | file = os.path.abspath(filename) if sys.platform == 'cygwin' else filename
|
40 | 44 | return os.path.normcase(os.path.realpath(os.path.normpath(file)))
|
| 45 | + |
| 46 | + |
| 47 | +@contextlib.contextmanager |
| 48 | +def paths_on_pythonpath(paths): |
| 49 | + """ |
| 50 | + Add the indicated paths to the head of the PYTHONPATH environment |
| 51 | + variable so that subprocesses will also see the packages at |
| 52 | + these paths. |
| 53 | +
|
| 54 | + Do this in a context that restores the value on exit. |
| 55 | +
|
| 56 | + >>> getfixture('monkeypatch').setenv('PYTHONPATH', 'anything') |
| 57 | + >>> with paths_on_pythonpath(['foo', 'bar']): |
| 58 | + ... assert 'foo' in os.environ['PYTHONPATH'] |
| 59 | + ... assert 'anything' in os.environ['PYTHONPATH'] |
| 60 | + >>> os.environ['PYTHONPATH'] |
| 61 | + 'anything' |
| 62 | +
|
| 63 | + >>> getfixture('monkeypatch').delenv('PYTHONPATH') |
| 64 | + >>> with paths_on_pythonpath(['foo', 'bar']): |
| 65 | + ... assert 'foo' in os.environ['PYTHONPATH'] |
| 66 | + >>> os.environ.get('PYTHONPATH') |
| 67 | + """ |
| 68 | + nothing = object() |
| 69 | + orig_pythonpath = os.environ.get('PYTHONPATH', nothing) |
| 70 | + current_pythonpath = os.environ.get('PYTHONPATH', '') |
| 71 | + try: |
| 72 | + prefix = os.pathsep.join(unique_everseen(paths)) |
| 73 | + to_join = filter(None, [prefix, current_pythonpath]) |
| 74 | + new_path = os.pathsep.join(to_join) |
| 75 | + if new_path: |
| 76 | + os.environ['PYTHONPATH'] = new_path |
| 77 | + yield |
| 78 | + finally: |
| 79 | + if orig_pythonpath is nothing: |
| 80 | + os.environ.pop('PYTHONPATH', None) |
| 81 | + else: |
| 82 | + os.environ['PYTHONPATH'] = orig_pythonpath |
0 commit comments