summaryrefslogtreecommitdiff
path: root/bitsandbytes/cuda_setup/env_vars.py
blob: 536a7d832e891f644023cebf5a46a987da76e062 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import os
from typing import Dict


def to_be_ignored(env_var: str, value: str) -> bool:
    ignorable = {
        "PWD",  # PWD: this is how the shell keeps track of the current working dir
        "OLDPWD",
        "SSH_AUTH_SOCK",  # SSH stuff, therefore unrelated
        "SSH_TTY",
        "HOME",  # Linux shell default
        "TMUX",  # Terminal Multiplexer
        "XDG_DATA_DIRS",  # XDG: Desktop environment stuff
        "XDG_RUNTIME_DIR",
        "MAIL",  # something related to emails
        "SHELL",  # binary for currently invoked shell
        "DBUS_SESSION_BUS_ADDRESS",  # hardware related
        "PATH",  # this is for finding binaries, not libraries
        "LESSOPEN",  # related to the `less` command
        "LESSCLOSE",
        "_",  # current Python interpreter
    }
    return env_var in ignorable


def might_contain_a_path(candidate: str) -> bool:
    return "/" in candidate


def is_active_conda_env(env_var: str) -> bool:
    return "CONDA_PREFIX" == env_var


def is_other_conda_env_var(env_var: str) -> bool:
    return "CONDA" in env_var


def is_relevant_candidate_env_var(env_var: str, value: str) -> bool:
    return is_active_conda_env(env_var) or (
        might_contain_a_path(value) and not
        is_other_conda_env_var(env_var) and not
        to_be_ignored(env_var, value)
    )


def get_potentially_lib_path_containing_env_vars() -> Dict[str, str]:
    return {
        env_var: value
        for env_var, value in os.environ.items()
        if is_relevant_candidate_env_var(env_var, value)
    }