OS Operations

Functions for working with the operating system.

---

EnvVarDictItemLower

Bases: TypedDict

Environment variable as a dictionary with lowercase keys.

EnvVarDictItemUpper

Bases: TypedDict

Environment variable as a dictionary with uppercase keys.

EnvVarFormat

Bases: Enum

Output format for environment variable conversions.

EnvVarItem dataclass

EnvVarItem(key, value)

Dataclass representing an environment variable key-value pair.

Attributes:

Name	Type	Description
key	str	The variable name.
value	str	The variable value.

from_any classmethod

from_any(value)

Create an EnvVarItem from a dict, tuple, or existing instance.

Parameters:

Name	Type	Description	Default
value	Any	An EnvVarItem, a 2-tuple, or a dict with key/Key and value/Value entries.	required

Returns:

Type	Description
EnvVarItem	A new EnvVarItem.

Raises:

Type	Description
ValueError	If value is not a supported type.

Source code in src/aibs_informatics_core/utils/os_operations.py

@classmethod
def from_any(cls, value: Any) -> "EnvVarItem":
    """Create an ``EnvVarItem`` from a dict, tuple, or existing instance.

    Args:
        value: An ``EnvVarItem``, a 2-tuple, or a dict with
            ``key``/``Key`` and ``value``/``Value`` entries.

    Returns:
        A new ``EnvVarItem``.

    Raises:
        ValueError: If ``value`` is not a supported type.
    """
    if isinstance(value, cls):
        return value
    elif isinstance(value, tuple):
        assert len(value) == 2
        return cls(*value)
    elif isinstance(value, dict):
        return cls(**{k.lower(): v for k, v in value.items()})
    else:
        raise ValueError(f"Invalid env_var type: {type(value)}")

to_dict

to_dict(
    lower: Literal[False] = False,
) -> EnvVarDictItemUpper

to_dict(lower: Literal[True]) -> EnvVarDictItemLower

to_dict(lower=False)

Convert to a TypedDict representation.

Parameters:

Name	Type	Description	Default
lower	bool	If True, use lowercase keys.	False

Returns:

Type	Description
EnvVarDictItemUpper | EnvVarDictItemLower	An EnvVarDictItemLower or EnvVarDictItemUpper.

Source code in src/aibs_informatics_core/utils/os_operations.py

def to_dict(self, lower: bool = False) -> EnvVarDictItemUpper | EnvVarDictItemLower:
    """Convert to a ``TypedDict`` representation.

    Args:
        lower: If True, use lowercase keys.

    Returns:
        An ``EnvVarDictItemLower`` or ``EnvVarDictItemUpper``.
    """
    if lower:
        return EnvVarDictItemLower(key=self.key, value=self.value)
    else:
        return EnvVarDictItemUpper(Key=self.key, Value=self.value)

to_tuple

to_tuple()

Convert to a (key, value) tuple.

Source code in src/aibs_informatics_core/utils/os_operations.py

def to_tuple(self) -> EnvVarTupleItem:
    """Convert to a ``(key, value)`` tuple."""
    return (self.key, self.value)

env_var_overrides

env_var_overrides(*env_vars)

Context manager to temporarily override environment variables.

On exit, all overridden variables are restored to their original values (or removed if they were unset).

Parameters:

Name	Type	Description	Default
*env_vars	EnvVarItemType	Environment variable items to set.	()

Source code in src/aibs_informatics_core/utils/os_operations.py

@contextmanager
def env_var_overrides(*env_vars: EnvVarItemType):
    """Context manager to temporarily override environment variables.

    On exit, all overridden variables are restored to their original values
    (or removed if they were unset).

    Args:
        *env_vars: Environment variable items to set.
    """
    original_env_vars = {}
    env_var_dict = to_env_var_dict(env_vars)  # type: ignore[arg-type]
    for key, value in env_var_dict.items():
        original_env_vars[key] = os.environ.get(key)
        os.environ[key] = value
    try:
        yield
    finally:
        for k, v in original_env_vars.items():
            if v is None:
                del os.environ[k]
            else:
                os.environ[k] = v

expandvars

expandvars(path, default=None, skip_escaped=False)

Expand environment variables of form $var and ${var}. If parameter 'skip_escaped' is True, all escaped variable references (i.e. preceded by backslashes) are skipped. Unknown variables are set to 'default'. If 'default' is None, they are left unchanged.

Source code in src/aibs_informatics_core/utils/os_operations.py

def expandvars(path, default=None, skip_escaped=False):
    """Expand environment variables of form $var and ${var}.
    If parameter 'skip_escaped' is True, all escaped variable references
    (i.e. preceded by backslashes) are skipped.
    Unknown variables are set to 'default'. If 'default' is None,
    they are left unchanged.
    """

    def replace_var(m):
        return os.environ.get(m.group(2) or m.group(1), m.group(0) if default is None else default)

    env_var_pattern = re.compile((r"(?<!\\)" if skip_escaped else "") + r"\$(\w+|\{([^}]*)\})")
    return re.sub(env_var_pattern, replace_var, path)

find_all_paths

find_all_paths(root, include_dirs=True, include_files=True)

Find all paths below root path

Parameters:

Name	Type	Description	Default
root	str | Path	root path to start	required
include_dirs	bool	Whether to include directories. Defaults to True.	True
include_files	bool	whether to include files. Defaults to True.	True

Returns: list of paths found under root

Source code in src/aibs_informatics_core/utils/os_operations.py

def find_all_paths(
    root: str | Path, include_dirs: bool = True, include_files: bool = True
) -> list[str]:
    """Find all paths below root path

    Args:
        root (str | Path): root path to start
        include_dirs (bool, optional): Whether to include directories. Defaults to True.
        include_files (bool, optional): whether to include files. Defaults to True.
    Returns:
        list of paths found under root
    """
    paths = []
    str_root = str(root) if isinstance(root, Path) else root
    if os.path.isfile(str_root) and include_files:
        paths.append(str_root)
    for parent, dirs, files in os.walk(str_root):
        if include_dirs:
            paths.extend([os.path.join(parent, name) for name in dirs])
        if include_files:
            paths.extend([os.path.join(parent, name) for name in files])
    return paths

generate_env_file_content

generate_env_file_content(env_vars)

Generate the content of an environment file.

Example

Given environment variables: - X = "${Y}/A" - Y = "${Z}/B" - Z = "C" Output: export Z=C export Y=${Z}/B export X=${Y}/A

Parameters:

Name	Type	Description	Default
env_vars	EnvVarCollection	A collection of environment variables.	required

Returns:

Type	Description
str	The content of the environment file.

Source code in src/aibs_informatics_core/utils/os_operations.py

def generate_env_file_content(env_vars: EnvVarCollection) -> str:
    """Generate the content of an environment file.

    Example:
        Given environment variables:
            - X = "${Y}/A"
            - Y = "${Z}/B"
            - Z = "C"
        Output:
            export Z=C
            export Y=${Z}/B
            export X=${Y}/A

    Args:
        env_vars (EnvVarCollection): A collection of environment variables.

    Returns:
        The content of the environment file.
    """
    return "\n".join(
        [
            f'export {k}="{v}"'
            # sort the remaining environment variables to proper order
            for k, v in order_env_vars(env_vars, env_var_format=EnvVarFormat.TUPLE)
        ]
    )

get_env_var

get_env_var(*keys: str) -> str | None

get_env_var(
    *keys: str, default_value: Literal[None]
) -> str | None

get_env_var(*keys: str, default_value: str) -> str

get_env_var(*keys, default_value=None)

get env variable using one of keys (sorted by priority)

Parameters:

Name	Type	Description	Default
keys	Tuple[str]	list of env keys to check (sorted based on fallback priority)	()
default_value	str | None	value to use if none are found. Defaults to none	None

Returns:

Type	Description
str | None	env variable value or default_value

Source code in src/aibs_informatics_core/utils/os_operations.py

def get_env_var(*keys: str, default_value: str | None = None) -> str | None:
    """get env variable using one of keys (sorted by priority)

    Arguments:
        keys (Tuple[str]): list of env keys to check
            (sorted based on fallback priority)
        default_value: value to use if none are found. Defaults to none

    Returns:
        env variable value or default_value
    """
    for key in keys:
        val = os.environ.get(key)
        if val:
            return val
    else:
        return default_value

order_env_vars

order_env_vars(
    env_vars: EnvVarCollection,
    env_var_format: Literal[OBJECT],
) -> list[EnvVarItem]

order_env_vars(
    env_vars: EnvVarCollection,
    env_var_format: Literal[TUPLE],
) -> list[EnvVarTupleItem]

order_env_vars(
    env_vars: EnvVarCollection,
    env_var_format: Literal[DICT_LOWER],
) -> list[EnvVarDictItemLower]

order_env_vars(
    env_vars: EnvVarCollection,
    env_var_format: Literal[DICT_UPPER],
) -> list[EnvVarDictItemUpper]

order_env_vars(env_vars, env_var_format=EnvVarFormat.TUPLE)

Resolve the order of environment variables based on their dependencies.

This function performs a topological sort on the environment variables, ordering them based on the dependencies found in their values. If a circular dependency is detected, a ValueError will be raised.

Parameters:

Name	Type	Description	Default
environment	Dict[str, str]	A dictionary where keys are environment variable names and values are their corresponding values. Dependencies are represented as ${var_name} or $var_name in the values.	required

Returns:

Type	Description
EnvVarCollection	List[Tuple[str, str]]: A list of tuples, where each tuple is a pair of environment variable name and its value, ordered based on their dependencies.

Raises:

Type	Description
ValueError	If a circular dependency is detected among the environment variables.

Source code in src/aibs_informatics_core/utils/os_operations.py

def order_env_vars(
    env_vars: EnvVarCollection, env_var_format: EnvVarFormat = EnvVarFormat.TUPLE
) -> EnvVarCollection:
    """Resolve the order of environment variables based on their dependencies.

    This function performs a topological sort on the environment variables,
    ordering them based on the dependencies found in their values. If a circular
    dependency is detected, a ValueError will be raised.

    Arguments:
        environment (Dict[str, str]): A dictionary where keys are environment variable names
            and values are their corresponding values. Dependencies are represented as
            ${var_name} or $var_name in the values.

    Returns:
        List[Tuple[str, str]]: A list of tuples, where each tuple is a pair of
            environment variable name and its value, ordered based on their dependencies.

    Raises:
        ValueError: If a circular dependency is detected among the environment variables.
    """
    # Create a list of environment keys
    env_var_dict = to_env_var_dict(env_vars)
    env_keys = list(env_var_dict.keys())

    # Initialize dependency maps
    env_key_deps: dict[str, set[str]] = {k: set() for k in env_keys}
    env_key_deps_rev: dict[str, set[str]] = {k: set() for k in env_keys}

    # Populate dependency maps
    for k, v in env_var_dict.items():
        for match in cast(list[tuple[str, str]], re.findall(r"(?:\$\{([^\}]+)\}|\$([\w]+))", v)):
            if (value := match[0] or match[1]) in env_keys:
                # Add dependencies
                env_key_deps[k].add(value)
                env_key_deps_rev[value].add(k)

    # Sort environment keys based on the number of dependencies
    unordered_env_keys = sorted(env_keys, key=lambda k: len(env_key_deps[k]), reverse=False)

    # Initialize the list of ordered environment key-value pairs
    ordered_env_key_pairs: list[EnvVarItem] = []

    # Order environment keys
    while unordered_env_keys:
        for i, k in enumerate(unordered_env_keys):
            # If a key has no dependencies, add it to the ordered list
            if not env_key_deps[k]:
                ordered_env_key_pairs.append(EnvVarItem(k, env_var_dict[k]))
                unordered_env_keys.pop(i)

                # Update dependency maps
                for rev_k in env_key_deps_rev.pop(k, set()):
                    env_key_deps[rev_k].discard(k)
                break
        else:
            # If a circular dependency is detected, raise an error
            raise ValueError("Circular dependency detected in environment variables")
    return to_env_var_list(ordered_env_key_pairs, env_var_format=env_var_format)  # type: ignore[call-overload]  # dynamic env var format doesnt work with mypy

set_env_var

set_env_var(key, value)

Set an environment variable.

Parameters:

Name	Type	Description	Default
key	str	The environment variable name.	required
value	str	The value to assign.	required

Source code in src/aibs_informatics_core/utils/os_operations.py

def set_env_var(key: str, value: str):
    """Set an environment variable.

    Args:
        key: The environment variable name.
        value: The value to assign.
    """
    os.environ[key] = value

to_env_var_dict

to_env_var_dict(env_vars)

Converts env vars to a dict

Parameters:

Name	Type	Description	Default
env_vars	Tuple[str, str] | Dict[str, str]	env vars to convert	required

Returns:

Type	Description
dict[str, str]	Dict[str, str]: dict of env vars

Source code in src/aibs_informatics_core/utils/os_operations.py

def to_env_var_dict(env_vars: EnvVarCollection) -> dict[str, str]:
    """Converts env vars to a dict

    Args:
        env_vars (Tuple[str, str] | Dict[str, str]): env vars to convert

    Returns:
        Dict[str, str]: dict of env vars
    """
    if isinstance(env_vars, dict):
        return env_vars

    env_var_dict = {}
    env_var_items = [EnvVarItem.from_any(env_var) for env_var in env_vars]
    for env_var_item in env_var_items:
        env_var_dict[env_var_item.key] = env_var_item.value
    return env_var_dict

to_env_var_list

to_env_var_list(
    env_vars: EnvVarCollection,
) -> list[EnvVarTupleItem]

to_env_var_list(
    env_vars: EnvVarCollection,
    env_var_format: Literal[TUPLE],
) -> list[EnvVarTupleItem]

to_env_var_list(
    env_vars: EnvVarCollection,
    env_var_format: Literal[OBJECT],
) -> list[EnvVarItem]

to_env_var_list(
    env_vars: EnvVarCollection,
    env_var_format: Literal[DICT_LOWER],
) -> list[EnvVarDictItemLower]

to_env_var_list(
    env_vars: EnvVarCollection,
    env_var_format: Literal[DICT_UPPER],
) -> list[EnvVarDictItemUpper]

to_env_var_list(
    env_vars, env_var_format=EnvVarFormat.TUPLE
)

Convert an environment variable collection to a list in the given format.

Parameters:

Name	Type	Description	Default
env_vars	EnvVarCollection	Environment variables as a dict or sequence.	required
env_var_format	EnvVarFormat	Desired output format.	TUPLE

Returns:

Type	Description
EnvVarSequence	A list of environment variables in the requested format.

Source code in src/aibs_informatics_core/utils/os_operations.py

def to_env_var_list(
    env_vars: EnvVarCollection, env_var_format: EnvVarFormat = EnvVarFormat.TUPLE
) -> EnvVarSequence:
    """Convert an environment variable collection to a list in the given format.

    Args:
        env_vars: Environment variables as a dict or sequence.
        env_var_format: Desired output format.

    Returns:
        A list of environment variables in the requested format.
    """
    if isinstance(env_vars, dict):
        env_var_list = [EnvVarItem(key, value) for key, value in env_vars.items()]
    else:
        env_var_list = [EnvVarItem.from_any(env_var) for env_var in env_vars]

    if env_var_format == EnvVarFormat.OBJECT:
        return env_var_list
    elif env_var_format == EnvVarFormat.TUPLE:
        env_var_tuple_list = [_.to_tuple() for _ in env_var_list]
        return env_var_tuple_list
    elif env_var_format == EnvVarFormat.DICT_LOWER:
        env_var_dict_list = [_.to_dict(lower=True) for _ in env_var_list]
        return env_var_dict_list
    elif env_var_format == EnvVarFormat.DICT_UPPER:
        env_var_dict_list = [_.to_dict(lower=False) for _ in env_var_list]
        return env_var_dict_list
    else:
        raise ValueError(f"Invalid env_var_format: {env_var_format}")

write_env_file

write_env_file(env_vars, path)

Write environment variables to a file.

Parameters:

Name	Type	Description	Default
env_vars	EnvVarCollection	A collection of environment variables.	required
path	Union[str, Path]	The path of the environment file.	required

Source code in src/aibs_informatics_core/utils/os_operations.py

def write_env_file(env_vars: EnvVarCollection, path: str | Path):
    """Write environment variables to a file.

    Args:
        env_vars (EnvVarCollection): A collection of environment variables.
        path (Union[str, Path]): The path of the environment file.
    """
    path = Path(path)
    path.write_text(generate_env_file_content(env_vars))