visualize disk_usage (mainly for Windows)

disk_usage_graph.py

# -*- coding: utf-8 -*-
# require: python 3.x, pillow, matplotlib, psutil
import sys, os, io
import time
from datetime import datetime
import pickle
import tkinter
import tkinter.ttk as ttk
from tkinter import E, W, N, S
import multiprocessing
import psutil
import matplotlib.pyplot as plt
from PIL import Image, ImageTk


__MYNAME__, _ = os.path.splitext(
    os.path.basename(sys.modules[__name__].__file__))
MAXHIST = 2 * 5 * 10000
MAXHISTSECS = 3 * 24 * 60 * 60
_GRAPH_IMGFN = ".{}_{{}}.png".format(__MYNAME__)
_savepath = os.path.join(
    os.environ.get("HOME", os.environ.get("USERPROFILE")),
    ".{}".format(__MYNAME__))
if not os.path.exists(_savepath):
    os.makedirs(_savepath)
_histpers = os.path.join(_savepath, "last.pickle")


def _imagefn(inidir):
    return os.path.join(_savepath, _GRAPH_IMGFN.format(inidir[0]))


def _nowfs(inidir):
    dt, r = datetime.now(), psutil.disk_usage(inidir)
    return dt, r.free / (1024.**3), r.free / r.total * 100


def _plt(inidir, fshist, q):
    fshist[inidir].append(_nowfs(inidir))
    first = fshist[inidir][0]
    time.sleep(0.25)
    try:
        cnt = 0
        while True:
            cnt += 1
            if cnt > 100:
                pickle.dump(fshist, io.open(_histpers, "wb"), protocol=-1)
                cnt = 0
            fshist[inidir].append(_nowfs(inidir))
            while len(fshist[inidir]) > MAXHIST or \
                  (fshist[inidir][-1][0] - fshist[inidir][0][0]).total_seconds() > MAXHISTSECS:
                fshist[inidir].pop(0)

            fig, ax1 = plt.subplots(tight_layout=True)
            fig.set_size_inches(16.53 / 4 * 4, 11.69 / 4 * 2)
            ax1.plot(
                [dt for dt, fs, pct in fshist[inidir]],
                [fs for dt, fs, pct in fshist[inidir]])
            ax1.grid(True)
            last = fshist[inidir][-1]
            ax1.set_title("[{}] free: {:.2f} % (at {}) -> {:.2f} % (at {})".format(
                inidir,
                first[2], first[0].strftime("%m-%d %I:%M%p"),
                last[2], last[0].strftime("%m-%d %I:%M%p")
            ))
            fig.autofmt_xdate()
            fig.savefig(_imagefn(inidir))
            plt.close(fig)
            qv = q.get()
            if not qv:
                break
    finally:
        pickle.dump(fshist, io.open(_histpers, "wb"), protocol=-1)


class App(object):
    def __init__(self, inidir):
        self._inidir, _ = os.path.splitdrive(os.path.abspath(inidir))
        self._root = tkinter.Tk()

        self._pctl = ttk.Progressbar(orient="vertical", maximum=100.)
        self._pctl.grid(row=0, column=0, sticky=N + S)
        self._imglbl = tkinter.Label(self._root)
        self._imglbl.grid(row=0, column=1)
        self._pctr = ttk.Progressbar(orient="vertical", maximum=100.)
        self._pctr.grid(row=0, column=2, sticky=N + S)
        #
        fshist = {}
        if os.path.exists(_histpers):
            fshist = pickle.load(io.open(_histpers, "rb"))
        if self._inidir not in fshist:
            fshist[self._inidir] = []
        fshist[self._inidir].append(_nowfs(self._inidir))
        self._pctl.config(value=fshist[self._inidir][0][-1])
        #
        self._q = multiprocessing.Queue()
        self._sp = multiprocessing.Process(
            target=_plt,
            args=(self._inidir, fshist, self._q))
        self._sp.start()
        #
        self._root.protocol("WM_DELETE_WINDOW", self._cleanup)
        #
        self._root.after(500, self._upd)
        #
        self._root.mainloop()

    def _upd(self, *args, **kwargs):
        self._q.put_nowait("_upd")
        self._pctr.config(value=_nowfs(self._inidir)[-1])
        try:
            img = Image.open(_imagefn(self._inidir))
            self._imgobj = ImageTk.PhotoImage(img)
            self._imglbl.config(image=self._imgobj)
        except Exception:
            pass
        self._root.after(5000, self._upd)

    def _cleanup(self):
        self._q.put(None)
        self._root.destroy()
        try:
            os.remove(_imagefn(self._inidir))
        except Exception:
            pass


if __name__ == '__main__':
    app = App(os.path.dirname(os.path.abspath(sys.argv[-1])))

splitdrive 以外は Windows 固有のことはしてないので、頑張れば Unix でも使えるとは思うけれど、Unix の場合はもっとシンプルなやり方はたくさんある、と思う。(そもそも「multiprocessing」してるのは Windows 版の matplotlib のある種の問題を避けるためである。)

rev=7 時点で、一定時間で「Fail to allocate bitmap」で落っこちるようになってしまってる。tk が出してるということは(tk86t.dll に対する strings で)確認したのだが、措置が全然わからない。しかも、先日までちっとも起こらず、昨日から突然起こるようになったんで、余計にわからない。