Skip to content

Commit

Permalink
feat: add --output option to list commands (supports text, json, csv …
Browse files Browse the repository at this point in the history
…and yaml)
  • Loading branch information
plaffitt committed Dec 27, 2024
1 parent 8a83370 commit b26fa90
Show file tree
Hide file tree
Showing 7 changed files with 86 additions and 18 deletions.
4 changes: 4 additions & 0 deletions src/pvecontrol/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,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):
Expand Down Expand Up @@ -75,6 +76,9 @@ 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", default=OutputFormats.TEXT, choices=[o.value for o in OutputFormats]
)
parser.add_argument(
"-c", "--cluster", action="store", required=True, help="Proxmox cluster name as defined in configuration"
)
Expand Down
4 changes: 2 additions & 2 deletions src/pvecontrol/actions/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions src/pvecontrol/actions/storage.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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)
5 changes: 3 additions & 2 deletions src/pvecontrol/actions/task.py
Original file line number Diff line number Diff line change
@@ -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,
)


Expand Down
4 changes: 2 additions & 2 deletions src/pvecontrol/actions/vm.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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)
54 changes: 44 additions & 10 deletions src/pvecontrol/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,13 @@
import sys
import re
import curses
import json

from collections import OrderedDict
from enum import Enum
from enum import Enum, StrEnum, auto

import yaml

from humanize import naturalsize
from prettytable import PrettyTable

Expand All @@ -20,6 +24,13 @@ class Fonts:
END = "\033[0m"


class OutputFormats(StrEnum):
TEXT = auto()
JSON = auto()
CSV = auto()
YAML = auto()


def terminal_support_colors():
try:
_stdscr = curses.initscr()
Expand Down Expand Up @@ -48,9 +59,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:
Expand All @@ -61,17 +70,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)
Expand All @@ -83,7 +113,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):
Expand All @@ -104,7 +138,7 @@ def print_taskstatus(task):
"user",
"starttime",
]
print_tableoutput([task], columns)
print_output([task], columns)


def print_task(proxmox, upid, follow=False, wait=False):
Expand Down Expand Up @@ -142,6 +176,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)
29 changes: 29 additions & 0 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit b26fa90

Please sign in to comment.