def test_save_to_json_create_file(self):
data = {}
try:
with tempfile.NamedTemporaryFile(mode='w+', suffix='.json') as tmpf:
filename = tmpf.name
save_to_json(data=data, filename=filename)
self.assertTrue(os.path.isfile(filename),
msg=f'{filename} was not created')
self.assertTrue(os.path.getsize(filename),
msg=f'{filename} is not empty')
except Exception as e:
self.fail(f'Exception occured when trying to create an empty JSON file {e}')
def test_save_to_json_full_data(self):
data = {
'test_01': [11, 22, 33],
'test_02': [44, 55, 66],
'test_03': [77, 88, 99]
}
try:
with tempfile.NamedTemporaryFile(mode='w+', suffix='.json') as tmpf:
save_to_json(data=data, filename=tmpf.name)
self.assertTrue(os.path.isfile(tmpf.name))
with open(tmpf.name) as f:
loaded = json.load(tmpf)
for key in data.keys():
for i, o in zip(loaded[key], data[key]):
self.assertEqual(o, i,
msg='Given two rows in a JSON files are not the same')
except Exception as e:
self.fail(f'Exception occured when trying to save a JSON file {e}')
class DiskUsage:
"""
Personalize and visualize the disk space usage in the terminal
"""
def __init__(self,
path: str = "",
exclude: list = None,
details: bool = False,
every: bool = False) -> None:
self.path = path
self.exclude = []
if exclude:
self.exclude = exclude
self.details = details
self.every = every
self._platform = platform.system(
) # Check on which platform vizex operates
def print_charts(self, options: Options = None) -> None:
"""
Prints the charts based on user selection type
Args:
options (Options): colors and symbols for printing
"""
if not options:
options = Options()
if self.path:
parts = self.grab_specific(self.path[0])
else:
parts = self.grab_partitions(exclude=self.exclude,
every=self.every)
chrt = HorizontalBarChart(options)
for partname in parts:
self.print_disk_chart(chrt, partname, parts[partname])
def print_disk_chart(self, ch: Chart, partname: str, part: dict) -> None:
"""Prints the disk data as a chart
Args:
ch (Chart): to print
partname (str): partition title
part (dict): parition data to be visualized
"""
pre_graph_text = self.create_stats(part)
footer = None
if self.details:
footer = self.create_details_text(part)
maximum = part["total"]
current = part["used"]
post_graph_text = create_usage_warning(part['percent'], 80, 60)
ch.chart(
post_graph_text=post_graph_text,
title=partname,
pre_graph_text=pre_graph_text,
footer=footer,
maximum=maximum,
current=current,
)
print()
def grab_root(self) -> dict:
"""
Grab the data for the root partition
return:
(dict) with column titles as keys
"""
return {
"total": psutil.disk_usage("/").total,
"used": psutil.disk_usage("/").used,
"free": psutil.disk_usage("/").free,
"percent": psutil.disk_usage("/").percent,
"fstype": psutil.disk_partitions(all=False)[0][2],
"mountpoint": "/",
}
def grab_partitions(self, exclude: list, every: bool) -> dict:
"""Grabs all the partitions data
Args:
exclude (list): of partitions to exclude
every (bool): if all the partitions should be grabbed
"""
disks = {}
# If we don't need every part we grab root seperately
if not every and self._platform != 'Windows':
disks['root'] = self.grab_root()
disk_parts = psutil.disk_partitions(all=every)
for disk in disk_parts[1:]:
# Exclude mounpoints created by snap
if disk.device.startswith("/dev/loop"):
continue
# Check that tmp is not slipping as partition
if disk.mountpoint.startswith("/tmp"):
continue
# Check that part name is not in the excluded list
if self.exclude and disk[1].split("/")[-1] in exclude:
continue
try:
if psutil.disk_usage(disk[1]).total > 0:
disks[disk[1].split("/")[-1]] = {
"total": psutil.disk_usage(disk[1]).total,
"used": psutil.disk_usage(disk[1]).used,
"free": psutil.disk_usage(disk[1]).free,
"percent": psutil.disk_usage(disk[1]).percent,
"fstype": disk.fstype,
"mountpoint": disk.mountpoint,
}
except Exception as e:
print(e)
continue
return disks
def grab_specific(self, disk_path: str) -> dict:
"""
Grabs data for the partition of the user specified path
Args:
disk_path (str): to the partition to grab
"""
disks = {}
disks[disk_path] = {
"total": psutil.disk_usage(disk_path).total,
"used": psutil.disk_usage(disk_path).used,
"free": psutil.disk_usage(disk_path).free,
"percent": psutil.disk_usage(disk_path).percent,
"fstype": "N/A",
"mountpoint": "N/A",
}
return disks
def create_details_text(self, disk: dict) -> str:
"""
Creates a string representation of a disk
Args:
disk (dict): text to print
"""
return f"fstype={disk['fstype']}\tmountpoint={disk['mountpoint']}"
def create_stats(self, disk: dict) -> str:
"""
Creates statistics as string for a disk
Args:
disk (dict): stats to print
"""
r = ints_to_human_readable(disk)
return f"Total: {r['total']}\t Used: {r['used']}\t Free: {r['free']}"
def save_data(self, filename: str) -> None:
"""
Outputs disks/partitions data as a CSV file
Args:
filename (str): for the saved file
"""
data = self.grab_partitions(self.exclude, self.every)
if (file_type := filename.split(".")[-1].lower()) == 'csv':
save_to_csv(data, filename)
elif file_type == 'json':
save_to_json(data, filename)