Disk usage change notification utility - Python

du-watch.py

#!/usr/bin/env python3

# Disk usage watcher:
# A script that reports changes to disk usage between runs (for mounted devices)
#
# Written for use on Linux systems.
#
# When the percentage change for a device is greater than USAGE_CHANGE_NOTIFY_THRESHOLD,
# a line will be printed for that device. If the change is not greater than the threshold,
# nothing will be printed for that device. Output looks like this:
#
#     $ ./du-watch.py
#     /media/scratch at 47.4% capacity (+3.08%)
#
# Disk usage is only updated in DATA_FILE if the change threshold is exceeded. This prevents
# small changes from adding up silently if the script is run frequently. That is, if the change
# threshold is set to 5%, a notification will ALWAYS be printed when the usage changes 5% or more
# since the last time a notification was printed.
#
# Command-line options:
#     -v    Output the current state for all mounts, without updating disk usage data in DATA_FILE.
#           This is useful to see the values being computed, without affecting the next normal run.
#
# Author: Stephen Holdaway <stephen@stecman.co.nz>
# Originally created: 10 June 2014
# Licence: MIT

import os
import sys
import json
import math
import time

from os.path import expanduser

DATA_FILE = expanduser("~/.duwatch")
USAGE_CHANGE_NOTIFY_THRESHOLD = 0.05
FORCE_DISPLAY = "-v" in sys.argv

def format_bytes(num):
	for x in ['B','KiB','MiB','GiB','TiB']:
		if num < 1024.0:
			return "%3.1f %s" % (num, x)
		num /= 1024.0

def get_previous_data():
	if os.path.exists(DATA_FILE):
		with open(DATA_FILE) as json_file:
			return json.load(json_file)
	else:
		return None

def update_data(devices, last_data={}):
	output = last_data
	for device in devices:
		output[device.mountPoint] = {
			"total": device.bytes_total,
			"used": device.bytes_used,
			"timestamp": time.time()
		}

	with open(DATA_FILE, "w") as fp:
		json.dump(output, fp, indent=True)

def get_mounted_devices():
	mounts = []

	with open("/proc/mounts") as f:
		for line in f:
			parts = line.split(" ")

			# Only use local block devices
			if parts[0][0:5] != "/dev/":
				continue

			mounts.append( Drive(parts[0], parts[1] ))

	return mounts


class Drive:
	def __init__(self, device, mountPoint):
		self.device = device
		self.mountPoint = mountPoint.replace("\\040", " ")

		self._stat = None

	@property
	def bytes_free(self):
		st = self._populate_space_stats()
		return st.f_bavail * st.f_frsize

	@property
	def bytes_total(self):
		st = self._populate_space_stats()
		return st.f_blocks * st.f_frsize

	@property
	def bytes_used(self):
		st = self._populate_space_stats()
		return (st.f_blocks - st.f_bfree) * st.f_frsize

	def _populate_space_stats(self):
		if not self._stat:
			self._stat = os.statvfs(self.mountPoint)

		return self._stat


devices = get_mounted_devices()
last_data = get_previous_data()
no_save_devices = []

if not last_data:
	print("Warning: no previous data to compare against")
	last_data = {}

for drive in devices:
	# Skip if the mount point didn't exist before
	if drive.mountPoint in last_data:
		last_drive = last_data[drive.mountPoint]
	else:
		print("Warning: New mount found at '%s'" % drive.mountPoint)
		continue

	# Skip out if the sizes don't match
	if drive.bytes_total != last_drive["total"]:
		print("Warning: Mount point '%s' has changed size from %s to %s" % (
			drive.mountPoint,
			format_bytes(last_drive["total"]),
			format_bytes(drive.bytes_total)
		))
		continue

	bytes_diff = drive.bytes_used - last_drive["used"]
	change = float(bytes_diff) / float(drive.bytes_total)

	# If the drive is at more than 90% capacity, notify more regularly
	thresh = USAGE_CHANGE_NOTIFY_THRESHOLD
	if float(drive.bytes_used) / float(drive.bytes_total) > 0.9 and change > 0.0:
		thresh = 0.01

	# If the space usage has changed more than the threshold, print details
	if FORCE_DISPLAY or abs(change) >= thresh:
		print("%s at %.1f%% capacity (%+.2f%%)" % (
			drive.mountPoint,
			(drive.bytes_used / drive.bytes_total) * 100,
			change * 100
		))
	else:
		no_save_devices.append(drive)

# Don't overwrite the data of drives that haven't changed more than the threshold
for device in no_save_devices:
	devices.remove(device)

if not FORCE_DISPLAY:
	update_data(devices, last_data)