Automatically save all matplotlib plots to a database

autosave_mpl_background.py

'''
Automatically saves all you matplotlib plot using pickle to a sqlite database
such that you can open them later, interactively.

Put this file somewhere on your PYTHONPATH and add this to your .bashrc

export PYTHONPATH="/path/to/directory"
export MPLBACKEND="module://autosave_mpl_background"
'''

import os
import sys

import datetime
import datetime
import lzma
import queue
import threading
import inspect
import io
import pickle
import sqlite3
import subprocess

from matplotlib.backend_bases import Gcf
import matplotlib.backends.backend_qt5agg as qt5agg

DBFILE = '/data/plots/plots2.db'

SCHEMA = '''create table if not exists plots(
    id integer primary key autoincrement,
    png blob,
    pickle_lzma blob,
    filename text,
    pid integer,
    filecontent_lzma blob,
    at timestamp
)
'''

INDEX = '''
create index if not exists idx_at on plots (at);
'''

INSERT = '''insert into plots(
    png, pickle_lzma, filename, pid, filecontent_lzma, at
) values (?, ?, ?, ?, ?, ?)
'''

q = queue.Queue()
def save_thread():
    conn = sqlite3.connect(DBFILE)
    conn.execute(SCHEMA)
    conn.execute(INDEX)
    conn.commit()
    conn.close()
    while threading.main_thread().is_alive() or not q.empty():
        try:
            png, dump, filename, pid, filecontent, now = q.get(timeout=1)
        except queue.Empty:
            continue
        dump = lzma.compress(dump)
        filecontent = lzma.compress(filecontent)
        conn = sqlite3.connect(DBFILE)
        conn.execute(INSERT, (png, dump, filename, pid, filecontent, now))
        conn.commit()
        conn.close()
        print(f'saved plot', file=sys.stderr)

t = threading.Thread(target=save_thread)
t.start()

def save_fig(fig):
    now = datetime.datetime.now()
    caller = os.path.abspath(inspect.stack()[-1].filename)
    with open(caller, 'rb') as f:
        filecontent = f.read()
    pid = os.getpid()
    try:
        dump = pickle.dumps(fig)
    except Exception as ex:
        print(repr(ex))
        print('Can\'t pickle MPL figure, not saving plot')
        return
    figpng = io.BytesIO()
    fig.savefig(figpng, format='png', dpi=50)
    figpng = figpng.getvalue()

    item = figpng, dump, caller, pid, filecontent, now
    q.put(item)


class FigureManagerMPLTelegramSend(qt5agg.FigureManager):
    def show(self):
        save_fig(self.canvas.figure)
        super().show()

@qt5agg._BackendQT5Agg.export
class _BackendMPLTelegramSend(qt5agg._BackendQT5Agg):
    FigureManager = FigureManagerMPLTelegramSend

plot_saved_plots.py

#!/usr/bin/python

import os
del os.environ['MPLBACKEND']
import sys
import sqlite3
import io
import lzma
import pickle
import datetime
from PyQt5.QtCore import QSize, Qt, QDate, QRectF
from PyQt5.QtGui import QIcon, QPixmap, QPainter
from PyQt5.QtWidgets import (
        QApplication, QMainWindow, QTreeWidget, QTreeWidgetItem, QListView,
        QCalendarWidget, QHBoxLayout, QWidget, QListWidget, QListWidgetItem
        )
import collections

DBFILE = '/data/plots/plots2.db'

class CalendarWidget(QCalendarWidget):
    def __init__(self):
        super().__init__()
        self.marked = collections.Counter()

    def mark(self, ymd):
        self.marked[ymd] += 1

    def paintCell(self, painter, rect, date):
        painter.setRenderHint(QPainter.Antialiasing, True)
        ymd = date.year(), date.month(), date.day()
        if ymd in self.marked:
            painter.save()
            painter.drawRect(rect)
            painter.setPen(Qt.blue)
            text = f'\n{date.day()}\n({self.marked[ymd]})'
            painter.drawText(
                    QRectF(rect),
                    Qt.AlignCenter,
                    text)
            painter.restore()
        else:
            QCalendarWidget.paintCell(self, painter, rect, date)

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.dbpath = DBFILE
        self.conn = sqlite3.connect(self.dbpath)
        self.cal = None
        #self.setCentralWidget(self.build_tree())
        self.item_cache = {}
        self.setCentralWidget(self.build_cal_and_lw())
        self.fill_lw()

    def on_lw_dbl(self, item):
        row_id = item.data(Qt.UserRole)
        cur = self.conn.cursor()
        cur.execute('select pickle_lzma from plots where id = ?', (row_id,))
        b, = cur.fetchone()
        cur.close()
        print('decompressing...')
        b = lzma.decompress(b)
        print('done')
        fig = pickle.loads(b)
        fig.show()

    def on_cal(self):
        #a = self.cal.selectedDate()
        self.fill_lw()

    def fill_lw(self):
        lw = self.lw
        lw.clear()
        cur = self.conn.cursor()
        if self.cal is None:
            cur.execute('select id, png, filename, pid, at from plots')
        else:
            d = self.cal.selectedDate()
            a = datetime.datetime(d.year(), d.month(), d.day())
            b = a + datetime.timedelta(1)
            cur.execute('select id, png, filename, pid, at from plots where at >= ? and at <= ?', (a, b))
        rows = reversed(cur.fetchall())
        for row in rows:
            row_id, row_png, row_fn, row_pid, row_at = row
            pixmap = QPixmap()
            pixmap.loadFromData(row_png)
            icon = QIcon(pixmap)
            text = os.path.basename(row_fn)
            text = ':'.join(row_at.split()[1].split(':')[0:2])
            item = QListWidgetItem(icon, text)
            item.setData(Qt.UserRole, row_id)
            lw.addItem(item)
        cur.close()

    def build_lw(self):
        lw = QListWidget()
        lw.setViewMode(QListView.IconMode)
        lw.setIconSize(QSize(130,130))
        lw.setResizeMode(QListView.Adjust)
        self.lw = lw
        lw.itemDoubleClicked.connect(self.on_lw_dbl)
        return lw

    def build_cal(self):
        cal = CalendarWidget()
        cal.selectionChanged.connect(self.on_cal)

        cur = self.conn.cursor()
        cur.execute('select at from plots')
        for row_at, in cur.fetchall():
            day = row_at.split()[0]
            y,m,d = map(int, day.split('-'))
            cal.mark((y,m,d))
        cur.close()

        self.cal = cal
        return cal

    def build_cal_and_lw(self):
        layout = QHBoxLayout()
        layout.addWidget(self.build_cal())
        layout.addWidget(self.build_lw())
        container = QWidget()
        container.setLayout(layout)
        return container


def main():
    app = QApplication(sys.argv)
    main_window = MainWindow()
    main_window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()