Merge pull request #56 from brianhlin/sw2628_pbs_pro_config

Fix python loading of blah.config

src/scripts/blah.py

@@ -1,18 +1,29 @@
 """Common functions for BLAH python scripts"""
 
 import os
-import subprocess
+from ConfigParser import RawConfigParser
+from io import StringIO
 
-def load_env(config_dir):
-    """Load blah.config into the environment"""
-    load_config_path = os.path.join(config_dir, 'blah_load_config.sh')
-    command = ['bash', '-c', 'source %s && env' % load_config_path]
-    try:
-        config_proc = subprocess.Popen(command, stdout=subprocess.PIPE)
-        config_out, _ = config_proc.communicate()
+class BlahConfigParser(RawConfigParser, object):
+
+    def __init__(self, path='/etc/blah.config'):
+        # RawConfigParser requires ini-style [section headers] but since
+        # blah.config is also used as a shell script we need to fake one
+        self.header = 'blahp'
+        with open(path) as f:
+            config = f.read()
+        vfile = StringIO(u'[%s]\n%s' % (self.header, config))
+
+        super(BlahConfigParser, self).__init__()
+        # TODO: readfp() is replaced by read_file() in Python 3.2+
+        self.readfp(vfile)
+
+    def items(self):
+        return super(BlahConfigParser, self).items(self.header)
+
+    def get(self, option):
+        # ConfigParser happily includes quotes in value strings, which we
+        # happily allow in /etc/blah.config. This causes failures when joining
+        # paths, for example.
+        return super(BlahConfigParser, self).get(self.header, option).strip('"\'')
 
-        for line in config_out.splitlines():
-            (key, _, val) = line.partition('=')
-            os.environ[key] = val
-    except IOError:
-        pass

src/scripts/pbs_status.py

@@ -43,7 +43,7 @@
 import pickle
 import csv
 
-sys.path.insert(0, os.path.dirname(__file__))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 import blah
 
 cache_timeout = 60
@@ -233,7 +233,7 @@ def qstat(jobid=""):
     starttime = time.time()
     log("Starting qstat.")
     command = (qstat_bin, '-f')
-    if os.environ.get('pbs_pro').lower() != 'yes':
+    if config.get('pbs_pro').lower() != 'yes':
         command += ('-1',) # -1 conflicts with -f in PBS Pro
     if jobid:
         command += (jobid,)
@@ -357,15 +357,17 @@ def get_qstat_location():
     global _qstat_location_cache
     if _qstat_location_cache != None:
         return _qstat_location_cache
+
     try:
-        cmd = os.path.join(os.environ['pbs_binpath'], 'qstat')
+        location = os.path.join(config.get('pbs_binpath'), 'qstat')
     except KeyError:
         cmd = 'which qstat'
-    child_stdout = os.popen(cmd)
-    output = child_stdout.read()
-    location = output.split("\n")[0].strip()
-    if child_stdout.close():
-        raise Exception("Unable to determine qstat location: %s" % output)
+        child_stdout = os.popen(cmd)
+        output = child_stdout.read()
+        location = output.split("\n")[0].strip()
+        if child_stdout.close():
+            raise Exception("Unable to determine qstat location: %s" % output)
+
     _qstat_location_cache = location
     return location
 
@@ -527,7 +529,8 @@ def main():
     jobid = jobid_arg.split("/")[-1].split(".")[0]
 
     config_dir = os.path.dirname(os.path.abspath(__file__))
-    blah.load_env(config_dir)
+    global config
+    config = blah.BlahConfigParser()
 
     log("Checking cache for jobid %s" % jobid)
     cache_contents = None

src/scripts/slurm_status.py

@@ -43,7 +43,7 @@
 import pickle
 import csv
 
-sys.path.insert(0, os.path.dirname(__file__))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 import blah
 
 cache_timeout = 60
@@ -336,14 +336,15 @@ def get_slurm_location(program):
     if _slurm_location_cache != None:
         return os.path.join(_slurm_location_cache, program)
     try:
-        cmd = os.path.join(os.environ['slurm_binpath'], program)
+        location = os.path.join(config.get('slurm_binpath'), program)
     except KeyError:
         cmd = 'which %s' % program
-    child_stdout = os.popen(cmd)
-    output = child_stdout.read()
-    location = output.split("\n")[0].strip()
-    if child_stdout.close():
-        raise Exception("Unable to determine scontrol location: %s" % output)
+        child_stdout = os.popen(cmd)
+        output = child_stdout.read()
+        location = output.split("\n")[0].strip()
+        if child_stdout.close():
+            raise Exception("Unable to determine scontrol location: %s" % output)
+
     _slurm_location_cache = os.path.dirname(location)
     return location