diff --git a/omas/examples/omas_dynamic_imas.py b/omas/examples/omas_dynamic_imas.py index 8c751d6c8..3c08c6b58 100644 --- a/omas/examples/omas_dynamic_imas.py +++ b/omas/examples/omas_dynamic_imas.py @@ -12,15 +12,17 @@ # set OMAS_DEBUG_TOPIC to see when data is loaded dynamically os.environ['OMAS_DEBUG_TOPIC'] = 'dynamic' +user = os.environ.get('USER', "default_user") + # generate some data and save it in IMAS ods = ODS().sample(ntimes=2) -ods.save('imas', os.environ['USER'], 'DIII-D', 1000, 0, new=True, verbose=True) +ods.save('imas', user, 'DIII-D', 1000, 0, new=True, verbose=True) # ODS.open() will keep the file descriptor open so that OMAS # can load in memory only the data when it is first requested # NOTE: one can use the `with` statement or open()/close() ods = ODS() -with ods.open('imas', os.environ['USER'], 'DIII-D', 1000, 0): +with ods.open('imas', user, 'DIII-D', 1000, 0): # data gets read from IMAS when first requested print(ods['equilibrium.time_slice.:.global_quantities.ip']) # then it is in memory @@ -37,10 +39,10 @@ print(ods.flat().keys()) # continue loading more data -with ods.open('imas', os.environ['USER'], 'DIII-D', 1000, 0): +with ods.open('imas', user, 'DIII-D', 1000, 0): print(ods['equilibrium.time']) # tell us what IMAS elements have data -with ods.open('imas', os.environ['USER'], 'DIII-D', 1000, 0): +with ods.open('imas', user, 'DIII-D', 1000, 0): print(ods.keys()) print(ods['equilibrium.time_slice.0.profiles_1d'].keys()) diff --git a/omas/examples/save_load_all.py b/omas/examples/save_load_all.py index 65d5d9e39..9f0ea810a 100644 --- a/omas/examples/save_load_all.py +++ b/omas/examples/save_load_all.py @@ -34,10 +34,10 @@ def through_omas_suite(ods=None, test_type=None, do_raise=False): os.environ['OMAS_DEBUG_TOPIC'] = test_type ods1 = globals()['through_omas_' + test_type](ods) difference = ods.diff(ods1) - if not chedifferenceck: - print('OMAS data got saved and loaded correctly') - else: - pprint(difference) + + print('OMAS data got saved and loaded correctly') + print("Diff:") + pprint(difference) else: os.environ['OMAS_DEBUG_TOPIC'] = '*' diff --git a/omas/machine_mappings/mast.py b/omas/machine_mappings/mast.py index 4cd9458fa..cb97b6335 100644 --- a/omas/machine_mappings/mast.py +++ b/omas/machine_mappings/mast.py @@ -13,12 +13,12 @@ def get_pyuda_client(server=None, port=None): if server is not None: pyuda.Client.server = server - elif not os.environ['UDA_HOST']: + elif not os.environ.get('UDA_HOST'): raise pyuda.UDAException('Must set UDA_HOST environmental variable') if port is not None: pyuda.Client.port = port - elif not os.environ['UDA_PORT']: + elif not os.environ.get('UDA_PORT'): raise pyuda.UDAException('Must set UDA_PORT environmental variable') return pyuda.Client() diff --git a/omas/omas_imas.py b/omas/omas_imas.py index 5069c178c..45e579370 100644 --- a/omas/omas_imas.py +++ b/omas/omas_imas.py @@ -305,7 +305,7 @@ def save_omas_imas(ods, user=None, machine=None, pulse=None, run=None, occurrenc # handle default values for user, machine, pulse, run, imas_version # it tries to re-use existing information if user is None: - user = ods.get('dataset_description.data_entry.user', os.environ['USER']) + user = ods.get('dataset_description.data_entry.user', os.environ.get('USER', 'default_user')) if machine is None: machine = ods.get('dataset_description.data_entry.machine', None) if pulse is None: @@ -685,7 +685,7 @@ def browse_imas( """ # if no users are specified, find all users if user is None: - user = glob.glob(user_imasdbdir.replace('/%s/' % os.environ['USER'], '/*/')) + user = glob.glob(user_imasdbdir.replace('/%s/' % os.environ.get('USER', 'default_user'), '/*/')) user = list(map(lambda x: x.split(os.sep)[-3], user)) elif isinstance(user, str): user = [user] @@ -694,7 +694,7 @@ def browse_imas( imasdb = {} for username in user: imasdb[username] = {} - imasdbdir = user_imasdbdir.replace('/%s/' % os.environ['USER'], '/%s/' % username).strip() + imasdbdir = user_imasdbdir.replace('/%s/' % os.environ.get('USER', 'default_user'), '/%s/' % username).strip() # find MDS+ datafiles files = list(recursive_glob('*datafile', imasdbdir)) @@ -942,7 +942,7 @@ def through_omas_imas(ods, method=['function', 'class_method'][1]): :return: ods """ - user = os.environ['USER'] + user = os.environ.get('USER', 'default_user') machine = 'ITER' pulse = 1 run = 0 diff --git a/omas/omas_machine.py b/omas/omas_machine.py index 1a38e89c2..b2d24c42a 100644 --- a/omas/omas_machine.py +++ b/omas/omas_machine.py @@ -148,7 +148,7 @@ def machine_to_omas(ods, machine, pulse, location, options={}, branch='', user_m return ods else: return resolve_mapped(ods, machine, pulse, mappings, location, idm, options_with_defaults, branch, cache=cache) - + def resolve_mapped(ods, machine, pulse, mappings, location, idm, options_with_defaults, branch, cache=None): """ Routine to resolve a mapping @@ -206,7 +206,11 @@ def resolve_mapped(ods, machine, pulse, mappings, location, idm, options_with_d # ENVIRONMENTAL VARIABLE elif 'ENVIRON' in mapped: - data0 = data = os.environ[mapped['ENVIRON'].format(**options_with_defaults)] + data0 = data = os.environ.get(mapped['ENVIRON'].format(**options_with_defaults)) + if data is None: + raise ValueError( + f'Environmental variable {mapped["ENVIRON"].format(**options_with_defaults)} is not defined' + ) # PYTHON elif 'PYTHON' in mapped: diff --git a/omas/omas_mongo.py b/omas/omas_mongo.py index 912fa994c..83a2d863b 100644 --- a/omas/omas_mongo.py +++ b/omas/omas_mongo.py @@ -2,6 +2,7 @@ ------- ''' +import pathlib # to start a mongodb server on the local workstation # mongod --dbpath $DIRECTORY_WHERE_TO_STORE_DATA @@ -154,9 +155,9 @@ def get_mongo_credentials(server='', database='', collection=''): server = server.split('@')[-1] up = {'user': 'omas_test', 'pass': 'omas_test'} config = {} - filename = os.environ['HOME'] + '/.omas/mongo_credentials' - if os.path.exists(filename): - with open(filename) as f: + filepath = pathlib.Path.home() / '/.omas/mongo_credentials' + if filepath.exists(): + with open(filepath) as f: config = json.loads(f.read()) if 'default' in config: up = config['default'] diff --git a/omas/omas_uda.py b/omas/omas_uda.py index 7776b3f10..65512cce9 100644 --- a/omas/omas_uda.py +++ b/omas/omas_uda.py @@ -57,12 +57,12 @@ def load_omas_uda( if server is not None: pyuda.Client.server = server - elif not os.environ['UDA_HOST']: + elif not os.environ.get('UDA_HOST'): raise pyuda.UDAException('Must set UDA_HOST environmental variable') if port is not None: pyuda.Client.port = port - elif not os.environ['UDA_PORT']: + elif not os.environ.get('UDA_PORT'): raise pyuda.UDAException('Must set UDA_PORT environmental variable') # set this to get pyuda metadata (maybe of interest for future use): diff --git a/omas/tests/failed_imports.py b/omas/tests/failed_imports.py index 7eb6c3eb4..53fbf9782 100644 --- a/omas/tests/failed_imports.py +++ b/omas/tests/failed_imports.py @@ -1,3 +1,5 @@ +import pathlib + from omas.omas_setup import omas_rcparams import os import warnings @@ -40,7 +42,8 @@ from botocore.exceptions import NoCredentialsError import boto3 - if not os.path.exists(os.environ.get('AWS_CONFIG_FILE', os.environ['HOME'] + '/.aws/config')): + aws_confi_path = pathlib.Path(os.environ.get('AWS_CONFIG_FILE', pathlib.Path.home() / '/.aws/config')) + if not aws_confi_path.exists(): raise RuntimeError('Missing AWS configuration file ~/.aws/config') failed_S3 = False except (ImportError, RuntimeError, NoCredentialsError) as _excp: diff --git a/omas/tests/warning_setup.py b/omas/tests/warning_setup.py index 0caf473d7..d7a2f15e6 100644 --- a/omas/tests/warning_setup.py +++ b/omas/tests/warning_setup.py @@ -3,7 +3,7 @@ import os hard_warnings = True -print('Setting up OMAS warnings for user {}'.format(os.environ['USER'])) +print('Setting up OMAS warnings for user {}'.format(os.environ.get('USER', 'default_user'))) def set_omas_warnings():