codekingpro/portable-devtools

1import sys
2import subprocess
3import os
4import tempfile
5from types import ModuleType
6from typing import Dict, List
7
8base_install_dir = (
9    tempfile.gettempdir()
10    + f"/worker-{os.environ.get('PYTEST_XDIST_WORKER', 'unknown')}"
11    + "/persistence_test_chromadb_versions"
12)
13
14
15def get_path_to_version_install(version: str) -> str:
16    return base_install_dir + "/" + version
17
18
19def switch_to_version(version: str, versioned_modules: List[str]) -> ModuleType:
20    module_name = "chromadb"
21    # Remove old version from sys.modules, except test modules
22    old_modules = {
23        n: m
24        for n, m in sys.modules.items()
25        if n == module_name
26        or (n.startswith(module_name + "."))
27        or n in versioned_modules
28        or (any(n.startswith(m + ".") for m in versioned_modules))
29    }
30    for n in old_modules:
31        del sys.modules[n]
32    # Load the target version and override the path to the installed version
33    # https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
34    sys.path.insert(0, get_path_to_version_install(version))
35    import chromadb
36
37    assert chromadb.__version__ == version
38    return chromadb
39
40
41def get_path_to_version_library(version: str) -> str:
42    return get_path_to_version_install(version) + "/chromadb/__init__.py"
43
44
45def install_version(version: str, dep_overrides: Dict[str, str]) -> None:
46    # Check if already installed
47    version_library = get_path_to_version_library(version)
48    if os.path.exists(version_library):
49        return
50    path = get_path_to_version_install(version)
51    install(f"chromadb=={version}", path, dep_overrides)
52
53
54def install(pkg: str, path: str, dep_overrides: Dict[str, str]) -> int:
55    os.makedirs(path, exist_ok=True)
56
57    # -q -q to suppress pip output to ERROR level
58    # https://pip.pypa.io/en/stable/cli/pip/#quiet
59    command = [sys.executable, "-m", "pip", "-q", "-q", "install", pkg]
60
61    for dep, operator_version in dep_overrides.items():
62        command.append(f"{dep}{operator_version}")
63
64    # Only add --no-binary=chroma-hnswlib if it's in the dependencies
65    if "chroma-hnswlib" in pkg or any("chroma-hnswlib" in dep for dep in dep_overrides):
66        command.append("--no-binary=chroma-hnswlib")
67
68    command.append(f"--target={path}")
69
70    print(f"Installing chromadb version {pkg} to {path}")
71    return subprocess.check_call(command)
72