1#!/usr/bin/env python3
 2
 3"""Read/write OS environment variables.
 4This file is used by core_env during initialisation so don't include settings
 5or use logging.
 6Don't use the traits.py module (although it would be useful) because this code
 7gets imported early in startup.
 8"""
 9
10import os
11from datetime import timedelta
12
13from chart.common.path import Path
14
15
16def get_env(name, datatype=str, default=None, allow_none=False):
17    """Interpret an entry from the OS environment.
18
19    Zero length strings are returned as None.
20    bool values are always either True or False (no None).
21    int values can be None or an integer.
22
23    Args:
24        `name` (str): Environment variable to read
25        `datatype` (enum:'int'|'path'|'str'|'bool'):
26            Type to enterpret var as.
27
28    """
29
30    env = os.environ.get(name)
31    # print('get env ' + name + ' ' + str(datatype) + ' got ' + str(env))
32
33    if allow_none is True and env == 'None':
34        return None
35
36    if datatype == 'bool' or datatype is bool:
37        if env is None or len(env) == 0:
38            return False
39
40        return env[0] == '1' or env[0].lower() == 't'
41
42    elif datatype == 'int' or datatype is int:
43        if env is None:
44            return default
45
46        else:
47            return int(env)
48
49    elif datatype == 'path' or datatype is Path:
50        return Path(env)
51
52    elif datatype == 'str' or datatype is str:
53        if env is None or len(env) == 0:
54            return default
55
56        else:
57            return env
58
59    elif datatype == 'timedelta' or datatype is timedelta:
60        if env is None:
61            return default
62
63        else:
64            from chart.common.xml import xml_to_timedelta
65            return xml_to_timedelta(env)
66
67    else:
68        raise ValueError('unknown type {typ}'.format(typ=datatype))
69
70
71def set_env(name, value):
72    """Write to environment.
73    `name` should take the form "{ENV_PREFIX}MY_VARIABLE_NAME".
74    """
75    from chart.project import settings
76
77    if isinstance(value, bool):
78        value = '1' if value is True else '0'
79
80    os.environ[name.format(ENV_PREFIX=settings.ENV_PREFIX)] = value