diff --git a/src/pvecontrol/__init__.py b/src/pvecontrol/__init__.py index 9cb7b0e..20646a0 100644 --- a/src/pvecontrol/__init__.py +++ b/src/pvecontrol/__init__.py @@ -11,6 +11,7 @@ from pvecontrol import actions, node, vm, task, storage from pvecontrol.cluster import PVECluster from pvecontrol.config import set_config +from pvecontrol.utils import OutputFormats def action_test(proxmox, _args): @@ -76,6 +77,14 @@ def _parser(): parser = argparse.ArgumentParser(description="Proxmox VE control cli.", epilog="Made with love by Enix.io") parser.add_argument("-v", "--verbose", action="store_true") parser.add_argument("--debug", action="store_true") + parser.add_argument( + "-o", + "--output", + action="store", + type=OutputFormats, + default=OutputFormats.TEXT, + choices=list(OutputFormats), + ) parser.add_argument( "-c", "--cluster", action="store", required=True, help="Proxmox cluster name as defined in configuration" ) diff --git a/src/pvecontrol/actions/node.py b/src/pvecontrol/actions/node.py index f361125..a14a862 100644 --- a/src/pvecontrol/actions/node.py +++ b/src/pvecontrol/actions/node.py @@ -2,12 +2,12 @@ from pvecontrol.node import NodeStatus from pvecontrol.vm import VmStatus -from pvecontrol.utils import print_tableoutput, print_task +from pvecontrol.utils import print_output, print_task def action_nodelist(proxmox, args): """List proxmox nodes in the cluster using proxmoxer api""" - print_tableoutput(proxmox.nodes, columns=args.columns, sortby=args.sort_by, filters=args.filter) + print_output(proxmox.nodes, columns=args.columns, sortby=args.sort_by, filters=args.filter, output=args.output) # pylint: disable=too-many-branches,too-many-statements diff --git a/src/pvecontrol/actions/storage.py b/src/pvecontrol/actions/storage.py index e9c136d..3fa5555 100644 --- a/src/pvecontrol/actions/storage.py +++ b/src/pvecontrol/actions/storage.py @@ -1,5 +1,5 @@ from pvecontrol.storage import StorageShared, COLUMNS -from pvecontrol.utils import print_tableoutput +from pvecontrol.utils import print_output def action_storagelist(proxmox, args): @@ -19,4 +19,4 @@ def action_storagelist(proxmox, args): for _id, storage in storages.items(): storage["nodes"] = ", ".join(storage["nodes"]) - print_tableoutput(storages.values(), COLUMNS, sortby=args.sort_by, filters=args.filter) + print_output(storages.values(), COLUMNS, sortby=args.sort_by, filters=args.filter, output=args.output) diff --git a/src/pvecontrol/actions/task.py b/src/pvecontrol/actions/task.py index e8f2cf2..405ee9a 100644 --- a/src/pvecontrol/actions/task.py +++ b/src/pvecontrol/actions/task.py @@ -1,12 +1,13 @@ -from pvecontrol.utils import print_task, print_tableoutput +from pvecontrol.utils import print_task, print_output def action_tasklist(proxmox, args): - print_tableoutput( + print_output( proxmox.tasks, columns=args.columns, sortby=args.sort_by, filters=args.filter, + output=args.output, ) diff --git a/src/pvecontrol/actions/vm.py b/src/pvecontrol/actions/vm.py index 85afd14..905e17c 100644 --- a/src/pvecontrol/actions/vm.py +++ b/src/pvecontrol/actions/vm.py @@ -1,7 +1,7 @@ import logging import sys -from pvecontrol.utils import print_task, print_tableoutput +from pvecontrol.utils import print_task, print_output def _get_vm(proxmox, vmid): @@ -64,4 +64,4 @@ def action_vmmigrate(proxmox, args): def action_vmlist(proxmox, args): """List VMs in the Proxmox Cluster""" vms = proxmox.vms() - print_tableoutput(vms, columns=args.columns, sortby=args.sort_by, filters=args.filter) + print_output(vms, columns=args.columns, sortby=args.sort_by, filters=args.filter, output=args.output) diff --git a/src/pvecontrol/utils.py b/src/pvecontrol/utils.py index 3c34f4c..b98fbcb 100644 --- a/src/pvecontrol/utils.py +++ b/src/pvecontrol/utils.py @@ -3,9 +3,13 @@ import sys import re import curses +import json from collections import OrderedDict from enum import Enum + +import yaml + from humanize import naturalsize from prettytable import PrettyTable @@ -20,6 +24,16 @@ class Fonts: END = "\033[0m" +class OutputFormats(Enum): + TEXT = "text" + JSON = "json" + CSV = "csv" + YAML = "yaml" + + def __str__(self): + return self.value + + def terminal_support_colors(): try: _stdscr = curses.initscr() @@ -48,9 +62,7 @@ def teminal_support_utf_8(): ] -# Pretty output a table from a table of dicts -# We assume all dicts have the same keys and are sorted by key -def print_tableoutput(table, columns=None, sortby=None, filters=None): +def render_output(table, columns=None, sortby=None, filters=None, output=OutputFormats.TEXT): if not columns: columns = [] if not filters: @@ -61,17 +73,38 @@ def print_tableoutput(table, columns=None, sortby=None, filters=None): else: table = [filter_keys(n.__dict__ if hasattr(n, "__dict__") else n, columns) for n in table] - do_sort = not sortby is None + x = prepare_prettytable(table, sortby, filters) + + if sortby is not None: + sortby = "sortby" + + if output == OutputFormats.TEXT: + return x.get_string(sortby=sortby, fields=columns) + if output == OutputFormats.CSV: + return x.get_csv_string(sortby=sortby, fields=columns) + if output in (OutputFormats.JSON, OutputFormats.YAML): + json_string = x.get_json_string(sortby=sortby, fields=columns) + data = json.loads(json_string)[1:] + if output == OutputFormats.JSON: + return json.dumps(data) + return yaml.dump(data) + + return None + + +def prepare_prettytable(table, sortby, filters): + do_sort = sortby is not None x = PrettyTable() x.align = "l" x.field_names = [*table[0].keys(), "sortby"] if do_sort else table[0].keys() for line in table: + for key in line: + if isinstance(line[key], Enum): + line[key] = str(line[key]) if do_sort: line["sortby"] = line[sortby] - if isinstance(line[sortby], Enum): - line["sortby"] = str(line[sortby]) for key in NATURALSIZE_KEYS: if key in line: line[key] = naturalsize(line[key], binary=True) @@ -83,7 +116,11 @@ def print_tableoutput(table, columns=None, sortby=None, filters=None): for line in table: x.add_row(line.values()) - print(x.get_string(sortby="sortby" if do_sort else None, fields=columns)) + return x + + +def print_output(table, columns=None, sortby=None, filters=None, output=OutputFormats.TEXT): + print(render_output(table, columns, sortby, filters, output)) def filter_keys(input_d, keys): @@ -104,7 +141,7 @@ def print_taskstatus(task): "user", "starttime", ] - print_tableoutput([task], columns) + print_output([task], columns) def print_task(proxmox, upid, follow=False, wait=False): @@ -142,6 +179,6 @@ def print_task(proxmox, upid, follow=False, wait=False): time.sleep(1) print("") elif not wait: - print_tableoutput([{"log output": task.decode_log()}]) + print_output([{"log output": task.decode_log()}]) print_taskstatus(task) diff --git a/src/tests/test_utils.py b/src/tests/test_utils.py new file mode 100644 index 0000000..431ccf7 --- /dev/null +++ b/src/tests/test_utils.py @@ -0,0 +1,29 @@ +import json +import csv + +from io import StringIO +from unittest.mock import Mock + +import yaml + +from pvecontrol.vm import PVEVm, COLUMNS +from pvecontrol.utils import render_output, OutputFormats + + +def test_render_output(): + api = Mock() + vms = [ + PVEVm(api, "pve-node-1", 100, "running"), + PVEVm(api, "pve-node-1", 101, "running"), + PVEVm(api, "pve-node-2", 102, "stopped"), + ] + + output_text = render_output(vms, columns=COLUMNS, output=OutputFormats.TEXT) + output_json = render_output(vms, columns=COLUMNS, output=OutputFormats.JSON) + output_csv = render_output(vms, columns=COLUMNS, output=OutputFormats.CSV) + output_yaml = render_output(vms, columns=COLUMNS, output=OutputFormats.YAML) + + assert output_text.split("\n")[0].replace("+", "").replace("-", "") == "" + assert len(json.loads(output_json)) == 3 + assert len(list(csv.DictReader(StringIO(output_csv)))) == 3 + assert len(yaml.safe_load(output_yaml)) == 3