1
import sys
2
from site import addsitedir
3
 
4
VENDOR_PATH = "/session/metadata/vendor"
5
PYTHON_MODULES_PATH = "/session/metadata/python_modules"
6
 
7
 
8
def setup_python_search_path(*, LEGACY_VENDOR_PATH, PROCESS_PTH_FILES):
9
    # adjustSysPath adds the session path, but it is immortalised by the memory snapshot. This
10
    # code runs irrespective of the memory snapshot.
11
    if VENDOR_PATH in sys.path and LEGACY_VENDOR_PATH:
12
        sys.path.remove(VENDOR_PATH)
13
 
14
    if PYTHON_MODULES_PATH in sys.path:
15
        sys.path.remove(PYTHON_MODULES_PATH)
16
 
17
    # Insert vendor path after system paths but before site-packages
18
    # System paths are typically: ['/session', '/lib/python312.zip', '/lib/python3.12', '/lib/python3.12/lib-dynload']
19
    # We want to insert before '/lib/python3.12/site-packages' and other site-packages
20
    #
21
    # We also need the session path to be before the vendor path, if we don't do so then a local
22
    # import will pick a module from the vendor path rather than the local path. We've got a test
23
    # that reproduces this (vendor_dir).
24
    for i, path in enumerate(sys.path):
25
        if "site-packages" in path:
26
            if LEGACY_VENDOR_PATH:
27
                sys.path.insert(i, VENDOR_PATH)
28
            sys.path.insert(i, PYTHON_MODULES_PATH)
29
            break
30
    else:
31
        # If no site-packages found, fail
32
        raise ValueError("No site-packages found in sys.path")
33
 
34
    if PROCESS_PTH_FILES:
35
        # addsitedir processes any .pth files in PYTHON_MODULES_PATH, allowing
36
        # packages to extend sys.path declaratively. It also re-adds the
37
        # directory to sys.path, but it was already inserted above so this is a
38
        # no-op for that purpose.
39
        addsitedir(PYTHON_MODULES_PATH)