Merge pull request #281 from gafusion/fix-env-var-usage

Fix env var usage

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())

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'] = '*'

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()

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

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:

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']

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):

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:

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():