Update the planet-cache.py admin utility

- switch to argparse
- fix a couple of minor 2-3 bugs
- slightly nicer logging of a couple of exceptions

Author: offbyone

code/planet-cache.py

@@ -8,6 +8,7 @@
 __license__ = "Python"
 
 
+import argparse
 import configparser
 import os
 import shelve
@@ -45,11 +46,10 @@ def usage_error(msg, *args):
 
 def print_keys(item, title):
     keys = item.keys()
-    keys.sort()
-    key_len = max([len(k) for k in keys])
+    key_len = max([len(k) for k in sorted(keys)])
 
     print(title + ":")
-    for key in keys:
+    for key in sorted(keys):
         if item.key_type(key) == item.DATE:
             value = time.strftime(planet.TIMEFMT_ISO, item[key])
         else:
@@ -65,87 +65,122 @@ def fit_str(string, length):
 
 
 if __name__ == "__main__":
-    cache_file = None
-    want_ids = 0
     ids = []
 
-    command = None
-
-    for arg in sys.argv[1:]:
-        if arg == "-h" or arg == "--help":
-            usage()
-        elif arg == "-C" or arg == "--channel":
-            if command is not None:
-                usage_error("Only one command option may be supplied")
-            command = "channel"
-        elif arg == "-L" or arg == "--list":
-            if command is not None:
-                usage_error("Only one command option may be supplied")
-            command = "list"
-        elif arg == "-K" or arg == "--keys":
-            if command is not None:
-                usage_error("Only one command option may be supplied")
-            command = "keys"
-        elif arg == "-I" or arg == "--item":
-            if command is not None:
-                usage_error("Only one command option may be supplied")
-            command = "item"
-            want_ids = 1
-        elif arg == "-H" or arg == "--hide":
-            if command is not None:
-                usage_error("Only one command option may be supplied")
-            command = "hide"
-            want_ids = 1
-        elif arg == "-U" or arg == "--unhide":
-            if command is not None:
-                usage_error("Only one command option may be supplied")
-            command = "unhide"
-            want_ids = 1
-        elif arg.startswith("-"):
-            usage_error("Unknown option:", arg)
-        elif cache_file is None:
-            cache_file = arg
-        elif want_ids:
-            ids.append(arg)
-        else:
-            usage_error("Unexpected extra argument:", arg)
-
-    if cache_file is None:
+    parser = argparse.ArgumentParser(
+        description="Examine and modify information in the Planet cache."
+    )
+    parser.add_argument(
+        "-C",
+        "--channel",
+        action="store_const",
+        const="channel",
+        dest="command",
+        help="Display known information on the channel",
+    )
+    parser.add_argument(
+        "-L",
+        "--list",
+        action="store_const",
+        const="list",
+        dest="command",
+        help="List items in the channel",
+    )
+    parser.add_argument(
+        "-K",
+        "--keys",
+        action="store_const",
+        const="keys",
+        dest="command",
+        help="List all keys found in channel items",
+    )
+    parser.add_argument(
+        "-I",
+        "--item",
+        action="store_const",
+        const="item",
+        dest="command",
+        help="Display known information about the item(s)",
+    )
+    parser.add_argument(
+        "-H",
+        "--hide",
+        action="store_const",
+        const="hide",
+        dest="command",
+        help="Mark the item(s) as hidden",
+    )
+    parser.add_argument(
+        "-U",
+        "--unhide",
+        action="store_const",
+        const="unhide",
+        dest="command",
+        help="Mark the item(s) as not hidden",
+    )
+    parser.add_argument("cache_file", help="Cache file to operate on")
+    parser.add_argument(
+        "item_ids",
+        nargs="*",
+        help="Item IDs to operate on when using item-related commands",
+    )
+
+    args = parser.parse_args()
+
+    # Check if more than one command option was supplied
+    if "command" not in args or args.command is None:
+        usage_error("One command option must be supplied.")
+    elif (
+        len(
+            {
+                key
+                for key, value in vars(args).items()
+                if key == "command" and value is not None
+            }
+        )
+        > 1
+    ):
+        usage_error("Only one command option may be supplied")
+
+    # Handle missing cache_file
+    if not args.cache_file:
         usage_error("Missing expected cache filename")
-    elif want_ids and not len(ids):
+
+    # Handle commands that require item IDs
+    if args.command in ["item", "hide", "unhide"] and not args.item_ids:
         usage_error("Missing expected entry ids")
 
     # Open the cache file directly to get the URL it represents
     try:
-        with shelve.open(cache_file, "r") as db:
-            url = db[b"url"].decode("utf-8")
-    except shelve.error as e:
-        print(f"{cache_file}: {e!s}", file=sys.stderr)
-        sys.exit(1)
+        with shelve.open(args.cache_file, "r") as db:
+            url = db["url"]
     except KeyError:
-        print(f"{cache_file}: Probably not a cache file", file=sys.stderr)
+        print(f"{args.cache_file}: Probably not a cache file", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        print(f"{args.cache_file}: {e!s}", file=sys.stderr)
         sys.exit(1)
 
     # Now do it the right way :-)
     my_planet = planet.Planet(configparser.ConfigParser())
-    my_planet.cache_directory = os.path.dirname(cache_file)
+    my_planet.cache_directory = os.path.dirname(args.cache_file)
     channel = planet.Channel(my_planet, url)
 
-    for item_id in ids:
+    for item_id in args.item_ids:
         if not channel.has_item(item_id):
             print(item_id + ": Not in channel", file=sys.stderr)
             sys.exit(1)
 
     # Do the user's bidding
-    if command == "channel":
+    if args.command == "channel":
         print_keys(channel, "Channel Keys")
 
-    elif command == "item":
-        for item_id in ids:
+    elif args.command == "item":
+        for item_id in args.item_ids:
             item = channel.get_item(item_id)
             print_keys(item, "Item Keys for %s" % item_id)
 
-    elif command == "list":
+    elif args.command == "list":
         print("Items in Channel:")
         for item in channel.items(hidden=True, sorted=True):
             print("    " + item.id)
@@ -155,7 +190,7 @@ def fit_str(string, length):
             if hasattr(item, "hidden"):
                 print("         (hidden)")
 
-    elif command == "keys":
+    elif args.command == "keys":
         keys = {}
         for item in channel.items():
             for key in item:
@@ -170,8 +205,8 @@ def fit_str(string, length):
 
         print("Use --item to output values of particular items.")
 
-    elif command == "hide":
-        for item_id in ids:
+    elif args.command == "hide":
+        for item_id in args.item_ids:
             item = channel.get_item(item_id)
             if hasattr(item, "hidden"):
                 print(item_id + ": Already hidden.")
@@ -181,8 +216,8 @@ def fit_str(string, length):
         channel.cache_write()
         print("Done.")
 
-    elif command == "unhide":
-        for item_id in ids:
+    elif args.command == "unhide":
+        for item_id in args.item_ids:
             item = channel.get_item(item_id)
             if hasattr(item, "hidden"):
                 del item.hidden