dmenu-maim

#!/usr/bin/env python3
# dmenu-maim
# dmenu wrapper for maim (X11) or grim (Wayland), supporting saving to a file
# and copying to the system clipboard, using the entire screen, a selection, or
# the currently focused window.
#
# Last updated 2021-11-20
#
# Dependencies (X11): dmenu, maim (and slop), xdotool, xclip
# Dependencies (Wayland): dmenu, grim, slurp, wl-clipboard

# Copyright (c) 2020-21 Pratyush Venkatakrishnan
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from subprocess import Popen, PIPE, DEVNULL, run, check_output
from datetime import datetime
from os import path, remove, getenv
from tempfile import NamedTemporaryFile
from configparser import ConfigParser, NoSectionError, NoOptionError
from shlex import split

if getenv("WAYLAND_DISPLAY"):
    server = "wayland"
else:
    server = "x11"

configdir = getenv("XDG_CONFIG_HOME", f"{getenv('HOME')}/.config")
config = ConfigParser()
config.read(f"{configdir}/dmenu-maim/config.ini")

# dmenu command line (this will be followed by the prompt)
try:
    DMENU = split(config.get(server, "dmenu"))
except (NoSectionError, NoOptionError):
    DMENU = ["dmenu", "-i", "-p"]

# xclip command line (this will take the image file as stdin)
try:
    XCLIP = split(config.get(server, "xclip"))
except (NoSectionError, NoOptionError):
    if server == "x11":
        XCLIP = ["xclip", "-selection", "clipboard", "-t", "image/png"]
    else:
        XCLIP = ["wl-copy", "-t", "image/png"]

# The maim command line needs to be configured manually; there is no easy and
# Pythonic way to abstract that out

# Character used to indicate end of options
DONE = ";"

# Available options for screenshot (at least one of 's' or 'c' must be
# provided)
OPTIONS = {
    "s": "Save to file",
    "c": "Copy to clipboard",
    "p": "Include mouse pointer",
}

# Available modes for screenshot (exactly one must be provided)
MODES = {
    "d": "Desktop",
    "s": "Selection",
    "w": "Window",
}

opts = ""
done = False

# Get options for screenshot
while not done:
    optstring = f"Options ({opts})" if opts else "Options"
    dmenu = Popen([*DMENU, optstring],
                  stdin=PIPE, stdout=PIPE, text=True)

    # Populate list of available options
    dmenu.stdin.write(f"{DONE}: Done\n")
    for o, desc in OPTIONS.items():
        if o not in opts:
            dmenu.stdin.write(f"{o}: {desc}\n")

    # Get stdout from dmenu process, strip whitespace, only pay attention to
    # characters before a colon (if present), and convert to lowercase
    result = dmenu.communicate()[0].strip().split(":")[0].lower()
    for o in OPTIONS.keys():
        if o in result:
            opts += o

    if DONE in result or not result:
        done = True

if "s" not in opts and "c" not in opts:
    exit(1)

# Get screenshot mode
dmenu = Popen([*DMENU, "Mode"],
              stdin=PIPE, stdout=PIPE, text=True)
for m, desc in MODES.items():
    dmenu.stdin.write(f"{m}: {desc}\n")

# Get first character of stdout and convert to lowercase
mode = dmenu.communicate()[0].strip()[0].lower()

if mode not in MODES:
    exit(1)

if "s" in opts:
    timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    outfile = open(path.expanduser(
        f"~/images/Screenshots/screenshot-{timestamp}.png"
    ), mode="wb")
else:
    outfile = NamedTemporaryFile(mode="wb", delete=False)

if server == "x11":
    cmdline = ["maim"]
    if "p" not in opts:
        # Hide mouse pointer
        cmdline.append("-u")
    if mode == "s":
        # Grab a selection instead of the entire desktop
        cmdline.append("-s")
    elif mode == "w":
        # Grab current window
        current_window = run(["xdotool", "getactivewindow"],
                             capture_output=True, text=True).stdout.strip()
        cmdline.extend(["-i", current_window])

    run(cmdline, stdin=DEVNULL, stdout=outfile, stderr=DEVNULL)
    outfile.close()
else:
    cmdline = ["grim"]
    if "p" in opts:
        # Show mouse pointer
        cmdline.append("-c")
    if mode == "s":
        # Grab a selection instead of the entire desktop
        selection = check_output(["slurp"], text=True).strip()
        cmdline.extend(["-g", selection])
    elif mode == "w":
        # Grab current window
        current_window = check_output(
            r"""swaymsg -t get_tree \
                | jq -j '.. | select(.type?) | select(.focused).rect
                    | "\(.x),\(.y) \(.width)x\(.height)"'
            """,
            text=True, shell=True).strip()
        cmdline.extend(["-g", current_window])

    outfile.close()
    cmdline.append(outfile.name)
    run(cmdline, stdin=DEVNULL, stdout=DEVNULL, stderr=DEVNULL)

if "c" in opts:
    with open(outfile.name) as infile:
        run(XCLIP, stdin=infile)

if "s" not in opts:
    remove(outfile.name)