dotfiles/bin/screenshot_wayland.py

499 lines
14 KiB
Python
Raw Normal View History

2022-11-23 01:05:55 -07:00
#!/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
2023-04-28 22:14:43 -06:00
# - grim on wayland or scrot on x11
# - slurp on wayland slop on x11
2022-11-23 01:05:55 -07:00
# - swappy
2022-11-25 01:40:01 -07:00
import argparse, os, sys, time, re, shutil, tarfile, tempfile, pillow_avif
2022-11-23 01:05:55 -07:00
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
2022-11-25 01:39:08 -07:00
RELATIVE_DIR = Path("Pictures/screenshots_wayland") # Default save directory
2023-04-28 22:14:43 -06:00
DIR = Path.home() / RELATIVE_DIR
2022-11-23 01:05:55 -07:00
2022-11-23 19:15:04 -07:00
DIMENSIONS_REGEX = "([0-9]+),([0-9]+) ([0-9]+)x([0-9]+)"
2023-04-28 22:14:43 -06:00
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"
2022-11-23 19:15:04 -07:00
EDIT_REGEX = "([0-9]+(_[0-9]{0,3})?)_edit_([0-9]+)" + EXTENSION_REGEX
2023-04-28 22:14:43 -06:00
DS_BG_LIGHT = "white"
DS_SHADOW_LIGHT = "black"
DS_BG_DARK = "#d3869b"
DS_SHADOW_DARK = "black"
DEFAULT_EDIT_QUALITY = 50
2023-04-28 22:14:43 -06:00
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}"
2022-11-23 01:05:55 -07:00
# 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
2023-04-28 22:14:43 -06:00
# Returns a free path to the last screenshot in DIR
2022-11-27 22:46:57 -07:00
# (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]
2022-11-23 19:15:04 -07:00
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}"
2023-04-28 22:14:43 -06:00
2022-11-23 01:05:55 -07:00
# Copies the image to the wayland clipboard
def copy_to_clipboard(pic: Path):
2023-04-28 22:14:43 -06:00
# 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,
)
2022-11-23 01:05:55 -07:00
# Interactively gets user to select an area of the screen
2023-04-28 22:14:43 -06:00
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)
2022-11-23 01:05:55 -07:00
# 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
2023-04-28 22:14:43 -06:00
width = img.size[0] + border_width * 2
height = img.size[1] + border_width * 2
2022-11-23 01:05:55 -07:00
shadow = Image.new(mode, img.size, color=shadow_color)
canvas = Image.new(mode, (width, height), color=back_color)
2023-04-28 22:14:43 -06:00
canvas.paste(shadow, box=(border_width, int(border_width * 1.16)))
2022-11-23 01:05:55 -07:00
canvas = canvas.filter(GaussianBlur(border_width // 4))
canvas.paste(img, box=(border_width, border_width))
2022-11-25 01:39:08 -07:00
canvas.save(save, optimize=True)
2022-11-23 01:05:55 -07:00
2023-04-28 22:14:43 -06:00
# 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
2022-11-23 01:05:55 -07:00
if exact_dims is not None:
2023-04-28 22:14:43 -06:00
run(["scrot", "-a", exact_dims.as_scrot(), "-F", path]).check_returncode()
2022-11-23 01:05:55 -07:00
else:
2023-04-28 22:14:43 -06:00
run(["scrot", "-F", path]).check_returncode()
2022-11-23 01:05:55 -07:00
# 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:
2022-11-23 01:05:55 -07:00
shutil.copyfile(save_path, args.file)
2023-04-28 22:14:43 -06:00
# Applies an edit to the latest screenshot
def edit_subcommand(args):
2022-11-25 01:39:08 -07:00
og_path, edit_path = get_edit_path(args.extension)
2022-11-23 22:41:22 -07:00
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)))
2022-11-25 01:39:08 -07:00
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)
2023-04-28 22:14:43 -06:00
2022-11-25 01:40:01 -07:00
# 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()
2023-04-28 22:14:43 -06:00
ext = f".{args.extension}"
2022-11-25 01:40:01 -07:00
count = len(pics)
sys.stdout.write(
2023-04-28 22:14:43 -06:00
f"Compressing {count} images into {ext} @ quality {args.quality}\n"
)
sys.stdout.write(f"Progress: 0/{count}")
2022-11-25 01:40:01 -07:00
for i, pic in enumerate(pics):
2023-04-28 22:14:43 -06:00
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))
2023-04-28 22:14:43 -06:00
sys.stdout.write(f"\rProgress: {i+1}/{count}" + " " * 40)
2022-11-25 01:40:01 -07:00
2023-04-28 22:14:43 -06:00
sys.stdout.write("\nDone!\n")
2022-11-25 01:40:01 -07:00
2023-04-28 22:14:43 -06:00
pics = [TMP_DIR / pic.with_suffix(ext) for pic in pics]
2022-11-25 01:40:01 -07:00
# Pack into tar file
with tarfile.open(args.tar_path, "w:gz") as tar:
for pic in pics:
tar.add(pic, arcname=pic.name)
2023-04-28 22:14:43 -06:00
2022-11-23 01:05:55 -07:00
# Provides a drawing editor for the latest image
2022-11-27 22:46:57 -07:00
def markup_subcommand(args):
2023-04-28 22:14:43 -06:00
og_path, edit_path = get_edit_path("png")
2022-11-23 01:05:55 -07:00
2022-11-27 22:46:57 -07:00
if args.show_latest:
print(og_path)
return
run(["swappy", "-f", og_path, "-o", edit_path]).check_returncode()
2022-11-23 01:05:55 -07:00
if args.clipboard:
2022-11-27 22:46:57 -07:00
copy_to_clipboard(edit_path)
if args.file is not None:
shutil.copyfile(edit_path, args.file)
2022-11-23 01:05:55 -07:00
2023-04-28 22:14:43 -06:00
2022-11-23 01:05:55 -07:00
# ===================================================================
# Parse args
# ===================================================================
def parser_dir(s: str):
if os.path.isdir(s):
return Path(s)
else:
raise NotADirectoryError(f"`{s}` is not a directory")
2023-04-28 22:14:43 -06:00
def parse_dimensions(s: str) -> ScreenshotDimensions:
return ScreenshotDimensions.from_string(s)
2022-11-23 19:15:04 -07:00
def parse_percent(s: str) -> int:
s = s.strip()
try:
2023-04-28 22:14:43 -06:00
return int(re.fullmatch("([0-9]+)%?", s)[1])
2022-11-23 19:15:04 -07:00
except IndexError:
raise Exception(f"{s} is not a valid percent. Does not match /[0-9]+%?/")
2023-04-28 22:14:43 -06:00
2022-11-23 19:15:04 -07:00
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]+/")
2022-11-23 01:05:55 -07:00
2023-04-28 22:14:43 -06:00
2022-11-25 01:40:01 -07:00
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")
2023-04-28 22:14:43 -06:00
2022-11-23 01:05:55 -07:00
parser = argparse.ArgumentParser(
2023-04-28 22:14:43 -06:00
prog="Screenshot Wayland v1.0.0", description="Take a screenshot on Sway"
)
2022-11-23 01:05:55 -07:00
parser.add_argument(
2023-04-28 22:14:43 -06:00
"-s",
"--screenshot-dir",
2022-11-23 01:05:55 -07:00
type=parser_dir,
metavar="<dir>",
2022-11-25 01:39:08 -07:00
help=f"Save and edit directory. Default: ~/{RELATIVE_DIR}",
2023-04-28 22:14:43 -06:00
)
2022-11-23 01:05:55 -07:00
# Subcommands ====
2023-04-28 22:14:43 -06:00
subcommands = parser.add_subparsers(dest="subcommand", required=True)
2022-11-23 01:05:55 -07:00
# Take ====
2023-04-28 22:14:43 -06:00
take_subcmd = subcommands.add_parser("take", help="Takes a screenshot")
2022-11-23 01:05:55 -07:00
2023-04-28 22:14:43 -06:00
take_common = argparse.ArgumentParser(add_help=False)
take_common.add_argument(
2023-04-28 22:14:43 -06:00
"-c",
"--clipboard",
action="store_true",
help="Save the screenshot to your clipboard",
)
take_common.add_argument(
2023-04-28 22:14:43 -06:00
"file", nargs="?", type=Path, help="Save the screenshot to this file name"
)
2022-11-23 01:05:55 -07:00
# Different possible screenshot regions
2023-04-28 22:14:43 -06:00
region = take_subcmd.add_subparsers(dest="region", required=True)
full = region.add_parser(
2023-04-28 22:14:43 -06:00
"full", parents=[take_common], help="Take a screenshot of the entire screen"
)
2022-11-23 01:05:55 -07:00
exact = region.add_parser(
2023-04-28 22:14:43 -06:00
"exact",
parents=[take_common],
help="Exact dimensions of screenshot: 'x,y width,height'",
)
select = region.add_parser(
2023-04-28 22:14:43 -06:00
"select",
parents=[take_common],
help="Use `slurp` or `slop` to select a region with your mouse",
)
2022-11-23 01:05:55 -07:00
exact.add_argument(
2023-04-28 22:14:43 -06:00
"dimensions",
2022-11-23 19:15:04 -07:00
type=parse_dimensions,
2022-11-23 01:05:55 -07:00
metavar="'N,N NxN'",
2023-04-28 22:14:43 -06:00
help="Exact dimensions of screenshot: 'x,y widthxheight'",
)
2022-11-23 01:05:55 -07:00
# Edit ====
edit_subcmd = subcommands.add_parser(
2023-04-28 22:14:43 -06:00
"edit", help="Apply an edit to the latest screenshot"
)
2022-11-23 01:05:55 -07:00
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-r",
"--rescale",
2022-11-23 19:15:04 -07:00
type=parse_percent,
metavar="<N%>",
2023-04-28 22:14:43 -06:00
help="Rescale the latest screenshot to <precent> of the original",
)
2022-11-23 01:05:55 -07:00
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-s",
"--size",
2022-11-23 19:15:04 -07:00
type=parse_size,
2022-11-23 01:05:55 -07:00
metavar="<dims>",
2023-04-28 22:14:43 -06:00
help="Set the dimensions of the latest screenshot",
)
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-c",
"--clipboard",
action="store_true",
help="Save the edited screenshot to your clipboard",
)
2022-11-23 01:05:55 -07:00
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-d",
"--drop-shadow",
action="store",
choices=["light", "dark"],
help="Apply a drop shadow with a light/dark background",
)
2022-11-23 01:05:55 -07:00
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-e",
"--extension",
metavar="<ext>",
2022-11-23 19:15:04 -07:00
type=str,
2022-11-25 01:39:08 -07:00
default=DEFAULT_EDIT_EXTENSION,
2023-04-28 22:14:43 -06:00
help="Change image extension and image type saved",
)
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-q",
"--quality",
type=parse_percent,
2022-11-25 01:39:08 -07:00
default=DEFAULT_EDIT_QUALITY,
2023-04-28 22:14:43 -06:00
metavar="<N%>",
help="Set quality of new image. [0, 100], higher means bigger file",
)
2022-11-23 22:41:22 -07:00
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"--overwrite",
action="store_true",
help="Overwrite original image with the edited image",
)
edit_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"file",
nargs="?",
type=Path,
2022-11-25 01:39:08 -07:00
help="Save the edited screenshot to this file name",
2023-04-28 22:14:43 -06:00
)
2022-11-25 01:40:01 -07:00
# Archive ====
archive_subcmd = subcommands.add_parser(
2023-04-28 22:14:43 -06:00
"archive", help="Backup the screenshots directory to a tar file"
)
2022-11-25 01:40:01 -07:00
archive_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-e",
"--extension",
2022-11-25 01:40:01 -07:00
type=str,
2022-11-25 13:12:21 -07:00
default=DEFAULT_EDIT_EXTENSION,
2022-11-25 01:40:01 -07:00
metavar="<ext>",
2023-04-28 22:14:43 -06:00
help="Change image extension and image type backed up",
)
2022-11-25 01:40:01 -07:00
archive_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-q",
"--quality",
2022-11-25 01:40:01 -07:00
type=parse_percent,
2022-11-25 13:12:21 -07:00
default=DEFAULT_EDIT_QUALITY,
2023-04-28 22:14:43 -06:00
metavar="<N%>",
help="Set quality of backed up images. [0, 100], higher means bigger file",
)
2022-11-25 01:40:01 -07:00
archive_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"which",
metavar="<which>",
choices=["all", "unedited"],
help="One of {all,unedited}. Backup only unedited images or all of them",
)
2022-11-25 01:40:01 -07:00
archive_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"tar_path",
metavar="<tar-path>",
2022-11-25 01:40:01 -07:00
type=parse_tar,
help="Destination tar file. Should be .{tar,tgz,tar.gz}",
2023-04-28 22:14:43 -06:00
)
2022-11-25 01:39:08 -07:00
2022-11-23 01:05:55 -07:00
# Markup ====
markup_subcmd = subcommands.add_parser(
2023-04-28 22:14:43 -06:00
"markup", help="Markup the latest screenshot in swappy"
)
2022-11-23 01:05:55 -07:00
markup_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-c",
"--clipboard",
action="store_true",
help="Save the screenshot to your clipboard",
)
2022-11-23 01:05:55 -07:00
markup_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"-s",
"--show-latest",
action="store_true",
help="Show the path to the latest image and exit",
)
2022-11-27 22:46:57 -07:00
markup_subcmd.add_argument(
2023-04-28 22:14:43 -06:00
"file",
nargs="?",
type=Path,
2022-11-27 22:46:57 -07:00
help="Save the edited screenshot to this file name",
2023-04-28 22:14:43 -06:00
)
2022-11-27 22:46:57 -07:00
2023-04-28 22:14:43 -06:00
args = parser.parse_args()
2022-11-23 01:05:55 -07:00
# Second layer of parser checks
if args.screenshot_dir is not None:
DIR = args.screenshot_dir
if not os.path.isdir(DIR) and os.path.exists(DIR):
print(f"Path `{DIR}` already exists and is not a directory")
exit(1)
elif not os.path.exists(DIR):
os.makedirs(DIR)
match args.subcommand:
case "take":
take_subcommand(args)
case "edit":
edit_subcommand(args)
2022-11-25 01:40:01 -07:00
case "archive":
archive_subcommand(args)
2022-11-23 01:05:55 -07:00
case "markup":
2022-11-27 22:46:57 -07:00
markup_subcommand(args)