From: Ben Steffen <bds@mit.edu>
Date: Mon, 9 Dec 2019 02:09:30 +0000 (-0500)
Subject: Update invirt-quota to Python 3
X-Git-Url: http://xvm.mit.edu/gitweb/invirt/packages/invirt-database.git/commitdiff_plain/b5afac5e5159505fa7efebadc01681c987d00c10

Update invirt-quota to Python 3
---

diff --git a/scripts/invirt-quota b/scripts/invirt-quota
index 1efdae6..7e81df4 100755
--- a/scripts/invirt-quota
+++ b/scripts/invirt-quota
@@ -1,67 +1,54 @@
-#!/usr/bin/env python
+#!/usr/bin/python3
 
 """
-invirt-quota allows an administrator to set memory, disk, and VM quotas 
+invirt-quota allows an administrator to set memory, disk, and VM quotas
 for an owner.  Invoking with only an owner name returns the current quotas for
 that owner.  Setting a parameter to -1 restores the default.
 """
 
-import sys
-import optparse
+import argparse
 
 from invirt import database
 
-def main(argv):
-    parser = optparse.OptionParser(usage = '%prog <owner> [options]',
-            description = __doc__.strip())
-    parser.add_option('-m', '--ram-total',
-            type = 'int',
-            dest = 'ramtotal',
-            help = 'set total concurrent RAM quota')
-    parser.add_option('-n', '--ram-single',
-            type = 'int',
-            dest = 'ramsingle',
-            help = 'set single VM RAM quota')
-    parser.add_option('-d', '--disk-total',
-            type = 'int',
-            dest = 'disktotal',
-            help = 'set total disk quota')
-    parser.add_option('-e', '--disk-single',
-            type = 'int',
-            dest = 'disksingle',
-            help = 'set single VM disk quota')
-    parser.add_option('-v', '--vms-total',
-            type = 'int',
-            dest = 'vmstotal',
-            help = 'set total VM quota')
-    parser.add_option('-w', '--vms-active',
-            type = 'int',
-            dest = 'vmsactive',
-            help = 'set active VM quota')
-    opts, args = parser.parse_args()
 
-    if len(args) != 1:
-        parser.print_help(sys.stderr)
-        return 1
-    owner_id = args[0]
+def main():
+    parser = argparse.ArgumentParser(description='Set memory, disk, and VM quotas')
+
+    parser.add_argument('-m', '--ram-total', type=int, dest='ramtotal',
+                        help='Set total concurrent RAM quota')
+    parser.add_argument('-n', '--ram-single', type=int, dest='ramsingle',
+                        help='Set single VM RAM quota')
+    parser.add_argument('-d', '--disk-total', type=int, dest='disktotal',
+                        help='Set total disk quota')
+    parser.add_argument('-e', '--disk-single', type=int, dest='disksingle',
+                        help='Set single disk quota')
+    parser.add_argument('-v', '--vms-total', type=int, dest='vmstotal',
+                        help='Set total VM quota')
+    parser.add_argument('-w', '--vms-active', type=int, dest='vmsactive',
+                        help='Set active VM quota')
+    parser.add_argument('owner')
+
+    args = parser.parse_args()
+
     database.connect()
     database.session.begin()
-    
-    owner = database.session.query(database.Owner).filter_by(owner_id=owner_id).first()
+
+    owner = database.session.query(database.Owner).filter_by(owner_id=args.owner).first()
     if owner is None:
-        owner = database.Owner(owner_id=owner_id)
+        owner = database.Owner(owner_id=args.owner)
 
     for resource, scope in [('ram',  'total'), ('ram',  'single'),
                             ('disk', 'total'), ('disk', 'single'),
                             ('vms',  'total'), ('vms',  'active')]:
-        val = getattr(opts, resource+scope)
+        val = getattr(args, resource + scope)
         if val is not None:
-            setattr(owner, resource+'_quota_'+scope, val if val >= 0 else None)
+            setattr(owner, resource + '_quota_' + scope, val if val >= 0 else None)
             database.session.add(owner)
 
     database.session.commit()
-    print owner
-    return 0
+
+    print(owner)
+
 
 if __name__ == '__main__':
-    sys.exit(main(sys.argv))
+    main()