Annotate after screenshot (install dependencies given the script)

annotate_screenshot.py

#!/usr/bin/env python3

# Dependencies
#sudo apt install python3-gi gir1.2-gtk-3.0 gnome-screenshot wl-clipboard
# put this script somewhere after making it an executable
# add a custom shortcut key to settings > keyboard > custom shortcuts


#!/usr/bin/env python3

import os
import subprocess
import tempfile
import shutil

import gi

gi.require_version("Gtk", "3.0")
gi.require_version("Gdk", "3.0")
from gi.repository import Gtk, Gdk, GdkPixbuf

import cairo


def capture_screenshot():
    fd, path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    try:
        result = subprocess.run(
            ["gnome-screenshot", "-a", "-f", path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if result.returncode != 0 or not os.path.exists(path):
            return None
        pixbuf = GdkPixbuf.Pixbuf.new_from_file(path)
        return pixbuf
    finally:
        if os.path.exists(path):
            os.remove(path)


class Annotator(Gtk.Window):
    def __init__(self, screenshot_pixbuf):
        super().__init__(title="Annotate Screenshot")
        self.set_default_size(screenshot_pixbuf.get_width(), screenshot_pixbuf.get_height())
        self.screenshot_pixbuf = screenshot_pixbuf
        self.result_pixbuf = None
        self.drawing = False
        self.last_x = 0
        self.last_y = 0

        self.surface = cairo.ImageSurface(
            cairo.FORMAT_ARGB32,
            self.screenshot_pixbuf.get_width(),
            self.screenshot_pixbuf.get_height(),
        )

        vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=0)
        self.add(vbox)

        self.drawing_area = Gtk.DrawingArea()
        self.drawing_area.set_size_request(
            self.screenshot_pixbuf.get_width(),
            self.screenshot_pixbuf.get_height(),
        )
        self.drawing_area.connect("draw", self.on_draw)
        self.drawing_area.add_events(
            Gdk.EventMask.BUTTON_PRESS_MASK
            | Gdk.EventMask.BUTTON_RELEASE_MASK
            | Gdk.EventMask.POINTER_MOTION_MASK
        )
        self.drawing_area.connect("button-press-event", self.on_button_press)
        self.drawing_area.connect("button-release-event", self.on_button_release)
        self.drawing_area.connect("motion-notify-event", self.on_motion)
        vbox.pack_start(self.drawing_area, True, True, 0)

        hbox = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=6)
        hbox.set_border_width(6)

        copy_button = Gtk.Button(label="Copy")
        copy_button.connect("clicked", self.on_copy_clicked)
        hbox.pack_end(copy_button, False, False, 0)

        cancel_button = Gtk.Button(label="Cancel")
        cancel_button.connect("clicked", self.on_cancel_clicked)
        hbox.pack_end(cancel_button, False, False, 0)

        vbox.pack_start(hbox, False, False, 0)

        self.connect("destroy", self.on_destroy)

    def on_draw(self, widget, cr):
        Gdk.cairo_set_source_pixbuf(cr, self.screenshot_pixbuf, 0, 0)
        cr.paint()
        cr.set_source_surface(self.surface, 0, 0)
        cr.paint()
        return False

    def on_button_press(self, widget, event):
        if event.button == 1:
            self.drawing = True
            self.last_x = event.x
            self.last_y = event.y
        return True

    def on_button_release(self, widget, event):
        if event.button == 1:
            self.drawing = False
        return True

    def on_motion(self, widget, event):
        if self.drawing:
            cr = cairo.Context(self.surface)
            cr.set_source_rgba(1.0, 0.0, 0.0, 1.0)
            cr.set_line_width(4.0)
            cr.move_to(self.last_x, self.last_y)
            cr.line_to(event.x, event.y)
            cr.stroke()
            self.last_x = event.x
            self.last_y = event.y
            self.drawing_area.queue_draw()
        return True

    def compose_result(self):
        width = self.screenshot_pixbuf.get_width()
        height = self.screenshot_pixbuf.get_height()
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)
        cr = cairo.Context(surface)
        Gdk.cairo_set_source_pixbuf(cr, self.screenshot_pixbuf, 0, 0)
        cr.paint()
        cr.set_source_surface(self.surface, 0, 0)
        cr.paint()
        pixbuf = Gdk.pixbuf_get_from_surface(surface, 0, 0, width, height)
        return pixbuf

    def on_copy_clicked(self, button):
        self.result_pixbuf = self.compose_result()
        Gtk.main_quit()

    def on_cancel_clicked(self, button):
        self.result_pixbuf = None
        Gtk.main_quit()

    def on_destroy(self, widget):
        if Gtk.main_level() > 0:
            Gtk.main_quit()


def open_annotation_ui(screenshot_pixbuf):
    win = Annotator(screenshot_pixbuf)
    win.show_all()
    Gtk.main()
    return win.result_pixbuf


def copy_to_clipboard(pixbuf):
    fd, path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    try:
        pixbuf.savev(path, "png", [], [])
        cmd = None
        if shutil.which("wl-copy") is not None:
            cmd = ["wl-copy", "--type", "image/png"]
        elif shutil.which("xclip") is not None:
            cmd = ["xclip", "-selection", "clipboard", "-t", "image/png", "-i"]

        if cmd is not None:
            with open(path, "rb") as f:
                subprocess.run(cmd, stdin=f)
        else:
            display = Gdk.Display.get_default()
            clipboard = Gtk.Clipboard.get_for_display(display, Gdk.SELECTION_CLIPBOARD)
            clipboard.set_image(pixbuf)
            clipboard.store()
    finally:
        if os.path.exists(path):
            os.remove(path)


def main():
    screenshot = capture_screenshot()
    if screenshot is None:
        return
    annotated = open_annotation_ui(screenshot)
    if annotated is not None:
        copy_to_clipboard(annotated)
    else:
        copy_to_clipboard(screenshot)


if __name__ == "__main__":
    main()