132 lines
3.6 KiB
Python
Executable File
132 lines
3.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import os
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
import time
|
|
|
|
# autopep8: off
|
|
import gi
|
|
gi.require_version("Notify", "0.7")
|
|
from gi.repository import Notify, GLib
|
|
# autopep8: on
|
|
|
|
|
|
class ScreenshotType(Enum):
|
|
FULL = "full"
|
|
AREA = "area"
|
|
WINDOW = "window"
|
|
|
|
|
|
SCREENSHOT_DIR = os.path.join(
|
|
os.environ.get("XDG_PICTURES_DIR", os.path.expanduser("~/Pictures")),
|
|
"Screenshots"
|
|
)
|
|
|
|
|
|
def wait_until_file_exists(filepath: str, timeout: int = 5):
|
|
"""Wait until a file exists or timeout."""
|
|
start_time = time.time()
|
|
while not os.path.isfile(filepath):
|
|
if time.time() - start_time > timeout:
|
|
return False
|
|
time.sleep(0.1)
|
|
return True
|
|
|
|
|
|
def take_screenshot(typeStr: str, filepath: str):
|
|
type = ScreenshotType(typeStr)
|
|
currentDesktop = os.environ.get("XDG_CURRENT_DESKTOP", "")
|
|
if "Hyprland" in currentDesktop:
|
|
cmd = {
|
|
ScreenshotType.FULL: f"hyprshot -z -m output -m active -o {SCREENSHOT_DIR} -f ", # since I only have one monitor
|
|
ScreenshotType.AREA: f"hyprshot -z -m region -o {SCREENSHOT_DIR} -f ",
|
|
ScreenshotType.WINDOW: f"hyprshot -z -m window -o {SCREENSHOT_DIR} -f ",
|
|
}
|
|
os.system(f"{cmd[type]}{filepath}")
|
|
wait_until_file_exists(filepath)
|
|
elif "niri" in currentDesktop:
|
|
cmd = {
|
|
ScreenshotType.FULL: f"niri msg action screenshot-screen",
|
|
ScreenshotType.AREA: f"niri msg action screenshot",
|
|
ScreenshotType.WINDOW: f"niri msg action screenshot-window",
|
|
}
|
|
niriScreenshotPath = SCREENSHOT_DIR + os.sep + ".niri_screenshot.png"
|
|
if os.path.isfile(niriScreenshotPath):
|
|
os.remove(niriScreenshotPath)
|
|
os.system(cmd[type])
|
|
wait_until_file_exists(niriScreenshotPath)
|
|
if os.path.isfile(niriScreenshotPath):
|
|
os.rename(niriScreenshotPath, filepath)
|
|
wait_until_file_exists(filepath)
|
|
else:
|
|
print("Unsupported desktop environment.")
|
|
exit(1)
|
|
|
|
|
|
def edit_screenshot(filepath: str):
|
|
os.system(f"gradia {SCREENSHOT_DIR}{os.sep}{filepath}")
|
|
|
|
|
|
def file_name(dir, prefix="screenshot", ext=".png"):
|
|
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
return f"{prefix}_{timestamp}{ext}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Take screenshots with hyprshot.")
|
|
parser.add_argument(
|
|
"type",
|
|
choices=[t.value for t in ScreenshotType],
|
|
help="Type of screenshot to take.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# file path
|
|
os.makedirs(SCREENSHOT_DIR, exist_ok=True)
|
|
filename = file_name(SCREENSHOT_DIR)
|
|
filepath = os.path.join(SCREENSHOT_DIR, filename)
|
|
|
|
# take screenshot
|
|
take_screenshot(args.type, filepath)
|
|
|
|
# check if successful
|
|
if not os.path.isfile(filepath):
|
|
print("Failed to take screenshot.")
|
|
exit(1)
|
|
|
|
# create loop instance
|
|
loop = GLib.MainLoop()
|
|
editing = False
|
|
|
|
# callback on default action (edit)
|
|
def edit_callback(n, action, user_data):
|
|
global editing
|
|
editing = True
|
|
edit_screenshot(filename)
|
|
n.close()
|
|
loop.quit()
|
|
|
|
# callback on close
|
|
def close_callback(n):
|
|
global editing
|
|
if not editing:
|
|
loop.quit()
|
|
|
|
# notification
|
|
Notify.init("Screenshot Utility")
|
|
n = Notify.Notification.new(
|
|
"Screenshot taken",
|
|
"Click to edit"
|
|
)
|
|
n.add_action(
|
|
"default",
|
|
"Default",
|
|
edit_callback,
|
|
None
|
|
)
|
|
n.connect("closed", close_callback)
|
|
n.show()
|
|
loop.run()
|