#!/usr/bin/env python3
import os
import inspect
import shlex
from datetime import datetime, date
from types import ModuleType
[docs]def str_to_date(s):
"""Get python datetime object from string.
Args:
s: a string
Returns:
datetime object or None
"""
v = None
try:
l = len(s)
if l == 8:
v = datetime.strptime(s, "%Y%m%d")
elif l == 10:
v = datetime.strptime(s, "%Y%m%d%H")
elif l == 12:
v = datetime.strptime(s, "%Y%m%d%H%M")
elif l == 14:
v = datetime.strptime(s, "%Y%m%d%H%M%S")
except:
v = None
return v
[docs]def date_to_str(d, format="%Y%m%d%H%M"):
"""Get string from python datetime object.
By default it converts to YYYYMMDDHHMM format unless
told otherwise by passing a different format
Args:
d: datetime object
Returns:
string in YYYYMMDDHHMM or shorter version of it
"""
v = d.strftime(format)
return v
[docs]def str_to_type(s, return_string=0):
"""Check if the string contains a float, int, boolean, datetime, or just regular string.
This will be used to automatically convert environment variables to data types
that are more convenient to work with. If you don't want this functionality,
pass return_string = 1
Args:
s: a string
return_string: Set to 1 to return the string itself
Set to 2 to return the string itself only for a datetime object
Returns:
a float, int, boolean, datetime, or the string itself when all else fails
"""
s = s.strip("\"'")
if return_string != 1:
if s.lower() in ["true", "yes", "yeah"]:
return True
if s.lower() in ["false", "no", "nope"]:
return False
if s in ["None", "null"]:
return None
v = str_to_date(s)
if v is not None:
if return_string == 2:
return s
return v
# int
try:
v = int(s)
return v
except:
pass
# float
try:
v = float(s)
return v
except:
pass
return s
[docs]def type_to_str(v):
"""Given a float/int/boolean/date or list of these types, gets a string
representing their values
Args:
v: a variable of the above types
Returns:
a string
"""
if isinstance(v, bool):
return "TRUE" if v else "FALSE"
elif isinstance(v, (int, float)):
pass
elif isinstance(v, date):
return date_to_str(v)
elif v is None:
return ""
return str(v)
[docs]def list_to_str(v, oneline=False):
"""Given a string or list of string, construct a string
to be used on right hand side of shell environement variables
Args:
v: a string/number, list of strings/numbers, or null string('')
Returns:
A string
"""
if isinstance(v, str):
return v
if isinstance(v, list):
v = [type_to_str(i) for i in v]
if oneline or len(v) <= 4:
shell_str = '( "' + '" "'.join(v) + '" )'
else:
shell_str = '( \\\n"' + '" \\\n"'.join(v) + '" \\\n)'
else:
shell_str = f"{type_to_str(v)}"
return shell_str
[docs]def str_to_list(v, return_string=0):
"""Given a string, construct a string or list of strings.
Basically does the reverse operation of `list_to_string`.
Args:
v: a string
Returns:
a string, list of strings or null string('')
"""
if not isinstance(v, str):
return v
v = v.strip()
if not v:
return None
if (v[0] == "(" and v[-1] == ")") or (v[0] == "[" and v[-1] == "]"):
v = v[1:-1]
v = v.replace(",", " ")
tokens = shlex.split(v)
lst = []
for itm in tokens:
itm = itm.strip()
if itm == "":
continue
# bash arrays could be stored with indices ([0]=hello ...)
if "=" in itm:
idx = itm.find("=")
itm = itm[idx + 1 :]
lst.append(str_to_type(itm, return_string))
return lst
return str_to_type(v, return_string)
[docs]def set_env_var(param, value):
"""Set an environment variable
Args:
param: the variable to set
value: either a string, list of strings or None
Returns:
None
"""
os.environ[param] = list_to_str(value)
[docs]def get_env_var(param):
"""Get the value of an environement variable
Args:
param: the environement variable
Returns:
Returns either a string, list of strings or None
"""
if not param in os.environ:
return None
value = os.environ[param]
return str_to_list(value)
[docs]def import_vars(dictionary=None, target_dict=None, env_vars=None):
"""Import all (or select few) environment/dictionary variables as python global
variables of the caller module. Call this function at the beginning of a function
that uses environment variables.
Note that for read-only environmental variables, calling this function once at the
beginning should be enough. However, if the variable is modified in the module it is
called from, the variable should be explicitly tagged as `global`, and then its value
should be exported back to the environment with a call to export_vars()
import_vars() # import all environment variables
global MY_VAR, MY_LIST_VAR
MY_PATH = "/path/to/somewhere"
MY_LIST_VAR.append("Hello")
export_vars() # these exports all global variables
There doesn't seem to an easier way of imitating the shell script doing way of things, which
assumes that everything is global unless specifically tagged local, while the opposite is true
for python.
Args:
dictionary: source dictionary (default=os.environ)
target_dict: target dictionary (default=caller module's globals())
env_vars: list of selected environement/dictionary variables to import, or None,
in which case all environment/dictionary variables are imported
Returns:
None
"""
if dictionary is None:
dictionary = os.environ
if target_dict is None:
target_dict = inspect.stack()[1][0].f_globals
if env_vars is None:
env_vars = dictionary
else:
env_vars = {k: dictionary[k] if k in dictionary else None for k in env_vars}
for k, v in env_vars.items():
# Don't replace variable with empty value
if not ((k in target_dict) and (v is None or v == "")):
target_dict[k] = str_to_list(v)
[docs]def export_vars(dictionary=None, source_dict=None, env_vars=None):
"""Export all (or select few) global variables of the caller module's
to either the environement/dictionary. Call this function at the end of
a function that updates environment variables.
Args:
dictionary: target dictionary to set (default=os.environ)
source_dict: source dictionary (default=caller modules globals())
env_vars: list of selected environement/dictionary variables to export, or None,
in which case all environment/dictionary variables are exported
Returns:
None
"""
if dictionary is None:
dictionary = os.environ
if source_dict is None:
source_dict = inspect.stack()[1][0].f_globals
if env_vars is None:
env_vars = source_dict
else:
env_vars = {k: source_dict[k] if k in source_dict else None for k in env_vars}
for k, v in env_vars.items():
# skip functions and other unlikely variable names
if callable(v):
continue
if isinstance(v, ModuleType):
continue
if not k or k[0] == "_":
continue
dictionary[k] = list_to_str(v)