Merge branch 'master' of github.com:MikeDacre/python-cluster

bin/my_queue

@@ -6,10 +6,8 @@ Print simple and fast queue information for only the current user.
 
         AUTHOR: Michael D Dacre, mike.dacre@gmail.com
   ORGANIZATION: Stanford University
-       LICENSE: MIT License, property of Stanford, use as you wish
-       VERSION: 0.1
        CREATED: 2015-12-18 12:43
- Last modified: 2016-02-23 09:25
+ Last modified: 2016-06-26 10:30
 
    DESCRIPTION: Uses slurmpy and pyslurm to check the job queue for only
                 one user's jobs. Produces a very simple display, for full
@@ -25,50 +23,25 @@ Print simple and fast queue information for only the current user.
 
 ============================================================================
 """
+from __future__ import print_function
 import sys
-import slurmy
+import argparse
+import cluster
 
 
-def main(running=False, queued=False, outputcount=False, outputlist=False):
-    """ Run everything """
-    # Check that the user isn't a dumbass
-    if running and queued:
-        sys.stderr.write("ERROR --> You can't specify both " +
-                         "'running' and 'queued'\n")
-        sys.exit(1)
-    if outputcount and outputlist:
-        sys.stderr.write("ERROR --> You can't specify both " +
-                         "'count' and 'list'\n")
+def get_arguments(argv=None):
+    """Parse command line arguments from argv and return an argparse parser.
 
-    # Create queue object
-    queue = slurmy.queue()
-    if running:
-        jobs = queue.running
-    elif queued:
-        jobs = queue.queued
-    else:
-        jobs = queue.queue
+    :argv:    A list of command line arguments.
+    :returns: An argparse ArgumentParser object, fully built.
 
-    # Print requested output
-    if outputcount:
-        sys.stdout.write("{0}\n".format(str(len(jobs))))
-    elif outputlist:
-        sys.stdout.write("{0}\n".format(
-            ' '.join(list([str(i) for i in jobs.keys()]))))
-    else:
-        for k, v in jobs.items():
-            try:
-                sys.stdout.write('{0}\t{1}\n'.format(
-                    k, str(v['name']).strip("b'")))
-            except IOError:
-                break
+    """
+    if not argv:
+        argv = sys.argv[1:]
 
-if __name__ == '__main__' and '__file__' in globals():
-    import argparse
-
-    f_class = argparse.RawDescriptionHelpFormatter
-    parser  = argparse.ArgumentParser(description=__doc__,
-                                      formatter_class=f_class)
+    parser  = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter)
 
     j = parser.add_argument_group('Choose jobs to show, default is all')
     j.add_argument('-r', '--running', action='store_true',
@@ -82,12 +55,50 @@ if __name__ == '__main__' and '__file__' in globals():
     o.add_argument('-l', '--list', action='store_true',
                    help="Print space separated list of job numbers")
 
-    args = parser.parse_args()
+    return parser
+
+
+def main(argv=None):
+    """Parse command line options and print queue information."""
+
+    parser = get_arguments(argv)
+    args   = parser.parse_args(argv)
 
-    # Run the script
-    main(running=args.running,
-         queued=args.queued,
-         outputcount=args.count,
-         outputlist=args.list)
+    if args.running and args.queued:
+        sys.stderr.write("ERROR --> You can't specify both " +
+                         "'running' and 'queued'\n")
+        return 1
+
+    if args.count and args.list:
+        sys.stderr.write("ERROR --> You can't specify both " +
+                         "'count' and 'list'\n")
+        return 2
+
+    # Create queue object
+    queue = cluster.Queue(user='self')
+    if args.running:
+        jobs = queue.running
+    elif args.queued:
+        jobs = queue.queued
+    else:
+        jobs = queue.jobs
 
-# vim:fenc=utf-8 tabstop=4 expandtab shiftwidth=4 softtabstop=4
+    # Print requested output
+    if args.count:
+        print("{0}".format(str(len(jobs))))
+    elif args.list:
+        print("{0}".format(
+            ' '.join(list([str(i) for i in jobs.keys()]))))
+    else:
+        for jid, job in jobs.items():
+            try:
+                print('{0}\t{1}'
+                      .format(jid, job.name))
+            except IOError:
+                break
+
+    # Done
+    return 0
+
+if __name__ == '__main__' and '__file__' in globals():
+    sys.exit(main())

cluster/__init__.py

@@ -5,9 +5,9 @@
          AUTHOR: Michael D Dacre, mike.dacre@gmail.com
    ORGANIZATION: Stanford University
         LICENSE: MIT License, property of Stanford, use as you wish
-        VERSION: 0.6.1
+        VERSION: 0.6.1b
         CREATED: 2015-12-11 22:19
-  Last modified: 2016-06-22 16:37
+  Last modified: 2016-07-15 21:23
  =============== ===================================================
 
 Allows simple job submission with *dependency tracking and queue waiting* with
@@ -142,6 +142,9 @@
 import signal as _signal
 import atexit as _atexit
 
+# Version Number
+__version__ = '0.6.1b'
+
 #################################################
 #  Currently configured job submission systems  #
 #################################################

cluster/queue.py

@@ -7,7 +7,7 @@
   ORGANIZATION: Stanford University
        LICENSE: MIT License, property of Stanford, use as you wish
        CREATED: 2015-12-11
- Last modified: 2016-06-25 13:35
+ Last modified: 2016-06-26 10:30
 
    DESCRIPTION: Provides a class to monitor the torque, slurm, or local
                 jobqueue queues with identical syntax.
@@ -40,6 +40,7 @@
 import sys
 import pwd      # Used to get usernames for queue
 import socket   # Used to get the hostname
+import getpass  # Used to get usernames for queue
 from time import time, sleep
 from subprocess import check_output, CalledProcessError
 
@@ -107,13 +108,14 @@ def __init__(self, user=None, qtype=None):
         # Get user ID as an int UID
         if user:
             if user == 'self' or user == 'current':
-                self.uid = pwd.getpwnam(os.environ['USER']).pw_uid
+                self.user = getpass.getuser()
+                self.uid  = pwd.getpwnam(self.user).pw_uid
             elif user == 'ALL':
                 self.user = None
             else:
-                if isinstance(user, int) or isinstance(user, str) \
-                        and user.isdigit():
-                    self.uid = pwd.getpwuid(int(user))
+                if isinstance(user, int) or (isinstance(user, str)
+                                             and user.isdigit()):
+                    self.uid  = int(user)
                 else:
                     self.uid = pwd.getpwnam(str(user)).pw_uid
         else:
@@ -536,7 +538,11 @@ def torque_queue_parser(user=None):
                 job_state = 'completed'
             logme.log('Job {} state: {}'.format(job_id, job_state),
                       'debug')
-            nds = xmljob.find('exec_host').text.split('+')
+            ndsx = xmljob.find('exec_host')
+            if ndsx:
+                nds = ndsx.text.split('+')
+            else:
+                nds = []
             nodes = []
             for node in nds:
                 if '-' in node:

setup.py

@@ -67,7 +67,7 @@ def run_tests(self):
 
 setup(
     name='python-cluster',
-    version='0.6.1',
+    version='0.6.1b',
     description='Submit functions and shell scripts to torque, slurm, ' +
                 'or local machines',
     long_description=long_description,