#!/usr/bin/env python3 # The all-in-one solution to taking and modifying screenshots on sWayland # # Additional dependencies: # - Python's pillow library # - pip install --user pillow-avif-plugin # Until PIL merges avif support # - grim on wayland or scrot on x11 # - slurp on wayland slop on x11 # - swappy import argparse, os, sys, time, re, shutil, tarfile, tempfile, pillow_avif from subprocess import run, Popen, PIPE, DEVNULL from PIL import Image from PIL.ImageFilter import GaussianBlur from pathlib import Path from typing import * # Global constants RELATIVE_DIR = Path("Pictures/screenshots_wayland") # Default save directory DIR = Path.home() / RELATIVE_DIR DIMENSIONS_REGEX = "([0-9]+),([0-9]+) ([0-9]+)x([0-9]+)" SLOP_REGEX = "([0-9]+)x([0-9]+)\+([0-9]+)\+([0-9]+)" EXTENSION_REGEX = "\.([A-z0-9]+)$" ORIGINAL_REGEX = "([0-9]+(_[0-9]{0,3})?)\.png" EDIT_REGEX = "([0-9]+(_[0-9]{0,3})?)_edit_([0-9]+)" + EXTENSION_REGEX DS_BG_LIGHT = "white" DS_SHADOW_LIGHT = "black" DS_BG_DARK = "#d3869b" DS_SHADOW_DARK = "black" DEFAULT_EDIT_QUALITY = 50 DEFAULT_EDIT_EXTENSION = "avif" class ScreenshotDimensions: def __init__(self, x, y, w, h): self.x = int(x) self.y = int(y) self.w = int(w) self.h = int(h) @classmethod def from_string(self, s: str): s = s.strip() if m := re.fullmatch(DIMENSIONS_REGEX, s): return self(m[1], m[2], m[3], m[4]) elif m := re.fullmatch(SLOP_REGEX, s): return self(m[3], m[4], m[1], m[2]) raise Exception(f"`{s}` does not match pattern {DIMENSIONS_REGEX}") def as_grim(self) -> str: return f"{self.x},{self.y} {self.w}x{self.h}" def as_scrot(self) -> str: return f"{self.x},{self.y},{self.w},{self.h}" # Returns a path to an unused screenshot in DIR def get_sceenshot_path() -> Path: t = time.time() p = DIR / f"{round(t)}.png" if os.path.isfile(p): return DIR / f"{str(t).replace('.', '_')}.png" else: return p # Returns a free path to the last screenshot in DIR # (original_path, new_edit_path) # # Example: # To edit the screenshot `1669235949.png` it'll return `1669235949_edit_0.png` # If `1669235949_edit_0.png` exists it returns `1669235949_edit_1.png`... # Raises IndexError if no screenshot was found in DIR def get_edit_path(ext: str) -> (Path, Path): pics = [f for f in os.listdir(DIR) if os.path.isfile(DIR / f)] originals = [p for p in pics if re.fullmatch(ORIGINAL_REGEX, p)] originals.sort() latest = re.split(EXTENSION_REGEX, originals[-1])[0] edits = [p for p in pics if re.fullmatch(f"{latest}_edit_[0-9]+\.[A-z0-9]+", p)] edits_nums = [re.fullmatch(EDIT_REGEX, e)[3] for e in edits] new_index = max([int(n) for n in edits_nums] + [0]) + 1 return DIR / originals[-1], DIR / f"{latest}_edit_{new_index}.{ext}" # Copies the image to the wayland clipboard def copy_to_clipboard(pic: Path): # TODO: xcompatability try: if "wayland" in os.environ["XDG_SESSION_TYPE"]: with open(pic, "r") as img: run(["wl-copy"], stdin=img, timeout=4) return except KeyError: pass run( ["xclip", "-selection", "clipboard", "-t", "image/png", "-i", str(pic)], timeout=4, ) # Interactively gets user to select an area of the screen def select_area() -> ScreenshotDimensions: try: if "wayland" in os.environ["XDG_SESSION_TYPE"]: slurp = run(["slurp"], text=True, stdout=PIPE) slurp.check_returncode() return ScreenshotDimensions.from_string(slurp.stdout) except KeyError: pass slop = run(["slop"], text=True, stdout=PIPE) slop.check_returncode() return ScreenshotDimensions.from_string(slop.stdout) # Adds drop-shadow to an image. If `img` and `save` are the same path, the image # is overwritten def add_drop_shadow(img: Path, save: Path, back_color: str, shadow_color: str): with Image.open(img) as img: mode = img.mode border_width = max(img.size) // 20 # 5% of larger dim width = img.size[0] + border_width * 2 height = img.size[1] + border_width * 2 shadow = Image.new(mode, img.size, color=shadow_color) canvas = Image.new(mode, (width, height), color=back_color) canvas.paste(shadow, box=(border_width, int(border_width * 1.16))) canvas = canvas.filter(GaussianBlur(border_width // 4)) canvas.paste(img, box=(border_width, border_width)) canvas.save(save, optimize=True) # Uses grim for wayland, scropt for x to take a screenshot. Default to full # screen without dims. exact_dims must match a slurp regex def take_screenshot(path, exact_dims: ScreenshotDimensions = None): try: if "wayland" in os.environ["XDG_SESSION_TYPE"]: if exact_dims is not None: run(["grim", "-g", exact_dims.as_grim(), path]).check_returncode() else: run(["grim", path]).check_returncode() return except KeyError: pass # x11 version if exact_dims is not None: run(["scrot", "-a", exact_dims.as_scrot(), "-F", path]).check_returncode() else: run(["scrot", "-F", path]).check_returncode() # Takes a screenshot as specified by args def take_subcommand(args): save_path = get_sceenshot_path() match args.region: case "full": take_screenshot(save_path, None) case "exact": take_screenshot(save_path, args.dimensions) case "select": take_screenshot(save_path, select_area()) case _: raise Exception(f"Region '{args.region}' not recognized") if args.clipboard: copy_to_clipboard(save_path) if args.file is not None: shutil.copyfile(save_path, args.file) # Applies an edit to the latest screenshot def edit_subcommand(args): og_path, edit_path = get_edit_path(args.extension) if args.overwrite: edit_path = og_path # Drop shadow match args.drop_shadow: case "light": add_drop_shadow(og_path, edit_path, DS_BG_LIGHT, DS_SHADOW_LIGHT) case "dark": add_drop_shadow(og_path, edit_path, DS_BG_DARK, DS_SHADOW_DARK) # Resize with Image.open(edit_path if args.drop_shadow else og_path) as img: if args.size: img = img.resize(args.size) elif args.rescale: img = img.reduce(int(1 / (args.rescale / 100))) img.save(edit_path, quality=args.quality, method=6, optimize=True) if args.clipboard: copy_to_clipboard(edit_path) if args.file is not None: shutil.copyfile(edit_path, args.file) # Saves `DIR` into a tar file def archive_subcommand(args): # Get the list of images to backup all_pics = [f for f in os.listdir(DIR) if os.path.isfile(DIR / f)] pics = [p for p in all_pics if re.fullmatch(ORIGINAL_REGEX, p)] if args.which == "all": pics += [p for p in all_pics if re.fullmatch(EDIT_REGEX, p)] pics = [Path(p) for p in pics] # Write compressed pictures to tmpdir with progress bar TMP_DIR = tempfile.mkdtemp() ext = f".{args.extension}" count = len(pics) sys.stdout.write( f"Compressing {count} images into {ext} @ quality {args.quality}\n" ) sys.stdout.write(f"Progress: 0/{count}") for i, pic in enumerate(pics): og = DIR / pic out = TMP_DIR / pic.with_suffix(ext) with Image.open(og) as img: img.save(out, quality=args.quality, method=6, optimize=True) stat = os.stat(og) os.utime(out, times=(stat.st_atime, stat.st_mtime)) sys.stdout.write(f"\rProgress: {i+1}/{count}" + " " * 40) sys.stdout.write("\nDone!\n") pics = [TMP_DIR / pic.with_suffix(ext) for pic in pics] # Pack into tar file with tarfile.open(args.tar_path, "w:gz") as tar: for pic in pics: tar.add(pic, arcname=pic.name) # Provides a drawing editor for the latest image def markup_subcommand(args): og_path, edit_path = get_edit_path("png") if args.show_latest: print(og_path) return run(["swappy", "-f", og_path, "-o", edit_path]).check_returncode() if args.clipboard: copy_to_clipboard(edit_path) if args.file is not None: shutil.copyfile(edit_path, args.file) # =================================================================== # Parse args # =================================================================== def parser_dir(s: str): if os.path.isdir(s): return Path(s) else: raise NotADirectoryError(f"`{s}` is not a directory") def parse_dimensions(s: str) -> ScreenshotDimensions: return ScreenshotDimensions.from_string(s) def parse_percent(s: str) -> int: s = s.strip() try: return int(re.fullmatch("([0-9]+)%?", s)[1]) except IndexError: raise Exception(f"{s} is not a valid percent. Does not match /[0-9]+%?/") def parse_size(s: str) -> (int, int): s = s.strip() try: m = re.fullmatch("([0-9]+)[x ]([0-9]+)", s) return int(m[1]), int(m[2]) except: raise Exception(f"{s} does not match size regex /[0-9]+[x ][0-9]+/") def parse_tar(s: str) -> Path: s = s.strip() if re.search("\.t(ar|gz|ar\.gz)$", s): return Path(s) else: raise Exception(f"{s} is not a path to a .{{tar,tgz,tar.gz}} file") parser = argparse.ArgumentParser( prog="Screenshot Wayland v1.0.0", description="Take a screenshot on Sway" ) parser.add_argument( "-s", "--screenshot-dir", type=parser_dir, metavar="