source: ntrip/trunk/BNC/scripts/helpScreenshots/shoot.py

#!/usr/bin/env python3
"""Screenshots for src/bnchelp.html: BNC's GUI on a virtual X display.

Start the display first (Xvfb 2600x1800 + openbox, see start_display.sh), then

    shoot.py [options] [job ...]        jobs from jobs.py, default: all
    shoot.py --list

Options:
    --user NAME     account for the live streams, replaces 'Example:Configs' in
                    the example configurations; the password is taken from the
                    credentials file (~/.netrc), never from the command line
    --work DIR      working directory (default /tmp/bnc_screenshots); BNC runs
                    in DIR/run/<job>, the images are written to DIR/out
    --install       copy the produced FigureNN.png files into src/IMG

The screenshots use Qt's Fusion style and an English locale; the main window
has a client size of 1200x850. The figure numbers are those of bnchelp.html.
Environment: BNC_SHOT_DISPLAY (default :99), BNC_SHOT_PREFIX (directory of
the unpacked packages if xdotool is not installed, see start_display.sh).
"""

import argparse
import glob
import os
import shutil
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
BNC_ROOT = os.path.dirname(os.path.dirname(HERE))
BNC = os.path.join(BNC_ROOT, "bnc")
DISPLAY = os.environ.get("BNC_SHOT_DISPLAY", ":99")
PREFIX = os.environ.get("BNC_SHOT_PREFIX", os.path.expanduser("~/.cache/bnc_screenshot_tools"))

if shutil.which("xdotool"):
    XDOTOOL, XLIBS = "xdotool", ""
else:
    XDOTOOL = os.path.join(PREFIX, "root/usr/bin/xdotool")
    XLIBS = os.path.join(PREFIX, "root/usr/lib/x86_64-linux-gnu")
XENV = dict(os.environ, DISPLAY=DISPLAY, LD_LIBRARY_PATH=XLIBS)
BNCENV = dict(os.environ, DISPLAY=DISPLAY, QT_STYLE_OVERRIDE="Fusion", LC_ALL="C.UTF-8",
              LANG="C.UTF-8", TZ="UTC")

WORK = "/tmp/bnc_screenshots"   # set by main()
OUT = os.path.join(WORK, "out")
USER = None


def xdo(*args):
    r = subprocess.run([XDOTOOL] + [str(a) for a in args], env=XENV, capture_output=True,
                       text=True)
    return r.stdout.strip()


def wait_window(title, timeout=30, count=1):
    """id of the (last) visible window whose title contains title; with
    count > 1 the list of ids once that many exist"""
    t0 = time.time()
    while time.time() - t0 < timeout:
        ids = xdo("search", "--onlyvisible", "--name", title).split()
        if len(ids) >= count:
            time.sleep(0.5)
            return ids[-1] if count == 1 else ids
        time.sleep(0.3)
    raise RuntimeError("window '%s' not found" % title)


def new_window(before, timeout=60):
    """first visible window with a title that was not in before"""
    for _ in range(timeout):
        new = set(xdo("search", "--onlyvisible", "--name", "").split()) - before
        new = [w for w in new if xdo("getwindowname", w)]
        if new:
            return new[0]
        time.sleep(1)
    raise RuntimeError("no new window")


def visible_windows():
    return set(xdo("search", "--onlyvisible", "--name", "").split())


def geometry(wid):
    out = xdo("getwindowgeometry", "--shell", wid)
    g = dict(l.split("=") for l in out.splitlines())
    return int(g["X"]), int(g["Y"]), int(g["WIDTH"]), int(g["HEIGHT"])


def place(wid, x, y, w=None, h=None):
    """move (and resize) a window; repeated until the window manager applied it"""
    time.sleep(1.0)
    for _ in range(3):
        if w:
            xdo("windowsize", "--sync", wid, w, h)
        xdo("windowmove", "--sync", wid, x, y)
        time.sleep(0.5)
        gx, gy, gw, gh = geometry(wid)
        if (not w or (gw, gh) == (w, h)) and abs(gx - x) < 30 and abs(gy - y) < 40:
            break


def click(wid, x, y, button=1):
    """click at x, y relative to the client area of window wid"""
    xdo("mousemove", "--window", wid, x, y)
    time.sleep(0.2)
    xdo("click", button)
    time.sleep(0.8)


def key(wid, *keys):
    xdo("windowactivate", "--sync", wid)
    for k in keys:
        xdo("key", "--clearmodifiers", k)
        time.sleep(0.2)


def typetext(wid, text):
    xdo("windowactivate", "--sync", wid)
    xdo("type", "--delay", "40", text)


def minimize(wid):
    xdo("windowminimize", wid)
    time.sleep(0.5)


def activate(wid):
    xdo("windowactivate", "--sync", wid)
    time.sleep(0.5)


def pointer_away():
    w, h = xdo("getdisplaygeometry").split()
    xdo("mousemove", int(w) - 10, int(h) - 10)
    time.sleep(0.5)


def saved(path):
    size = subprocess.run(["identify", "-format", "%wx%h", path], capture_output=True,
                          text=True).stdout
    print("  saved", os.path.relpath(path, WORK), size)
    return path


def frame():
    """screen rectangle of everything shown (windows incl. title bars) =
    bounding box of the non-black part of the root window"""
    raw = os.path.join(OUT, "_bbox.png")
    subprocess.run(["import", "-window", "root", raw], env=XENV, check=True)
    bb = subprocess.run(["convert", raw, "-format", "%@", "info:"], capture_output=True,
                        text=True).stdout.strip()            # WxH+X+Y
    wh, x, y = bb.split("+")
    w, h = wh.split("x")
    return int(x), int(y), int(w), int(h)


def shot(name, crop=None):
    """screenshot of all windows shown, or of crop=(x, y, w, h) in screen
    coordinates (e.g. from frame()); saved as OUT/<name>.png"""
    pointer_away()
    raw = os.path.join(OUT, "_root.png")
    subprocess.run(["import", "-window", "root", raw], env=XENV, check=True)
    path = os.path.join(OUT, name + ".png")
    if crop:
        x, y, w, h = crop
        subprocess.run(["convert", raw, "-crop", "%dx%d+%d+%d" % (w, h, x, y), "+repage", path],
                       check=True)
    else:
        subprocess.run(["convert", raw, "-trim", "+repage", path], check=True)
    return saved(path)


def wshot(name, wid):
    """screenshot of one window incl. its window manager frame"""
    pointer_away()
    path = os.path.join(OUT, name + ".png")
    subprocess.run(["import", "-frame", "-window", wid, path], env=XENV, check=True)
    return saved(path)


def crop_file(src, name, geometry):
    """crop an image file, geometry 'WxH+X+Y'"""
    path = os.path.join(OUT, name + ".png")
    subprocess.run(["convert", src, "-crop", geometry, "+repage", path], check=True)
    return saved(path)


class Bnc:
    """BNC instance in its own working directory (Input link, Output dir),
    started with a copy of an example configuration"""

    def __init__(self, job, conf=None, keys=(), live=False):
        self.dir = os.path.join(WORK, "run", job)
        shutil.rmtree(self.dir, ignore_errors=True)
        os.makedirs(os.path.join(self.dir, "Output"))
        os.symlink(os.path.join(BNC_ROOT, "Example_Configs", "Input"), os.path.join(self.dir, "Input"))
        cmd = [BNC]
        if conf:
            text = open(os.path.join(BNC_ROOT, "Example_Configs", conf), errors="replace").read()
            if live and USER:
                text = text.replace("Example:Configs@", USER + "@")
            open(os.path.join(self.dir, conf), "w").write(text)
            cmd += ["--conf", conf]
        for k, v in keys:
            if live and USER:
                v = str(v).replace("Example:Configs@", USER + "@")
            cmd += ["-key", k, str(v)]
        self.out = open(os.path.join(self.dir, "bnc.out"), "w")
        self.p = subprocess.Popen(cmd, cwd=self.dir, env=BNCENV, stdout=self.out,
                                  stderr=subprocess.STDOUT)
        self.main = wait_window("BKG Ntrip Client")

    def stop(self):
        if self.p.poll() is None:
            self.p.terminate()
            try:
                self.p.wait(10)
            except subprocess.TimeoutExpired:
                self.p.kill()
        time.sleep(1)


def cleanup():
    """end leftover BNC instances of this driver: processes with a window on
    the virtual display that run BNC_ROOT/bnc with this DISPLAY (nothing else)"""
    pids = set()
    for w in xdo("search", "--name", "").split():
        p = xdo("getwindowpid", w)
        if p:
            pids.add(int(p))
    for p in pids:
        try:
            cmd = open("/proc/%d/cmdline" % p, "rb").read().split(b"\0")[0].decode()
            env = open("/proc/%d/environ" % p, "rb").read().split(b"\0")
        except OSError:
            continue
        if cmd == BNC and ("DISPLAY=" + DISPLAY).encode() in env:
            os.kill(p, 15)
            print("  ended leftover BNC", p)
    time.sleep(1)


def main():
    global WORK, OUT, USER
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("jobs", nargs="*")
    ap.add_argument("--user")
    ap.add_argument("--work", default=WORK)
    ap.add_argument("--install", action="store_true")
    ap.add_argument("--list", action="store_true")
    args = ap.parse_args()

    # jobs.py imports this file as module 'shoot' (a second instance besides
    # __main__): the settings have to be made there
    sys.path.insert(0, HERE)
    import shoot as S
    S.WORK, S.USER = os.path.abspath(args.work), args.user
    S.OUT = os.path.join(S.WORK, "out")
    os.makedirs(S.OUT, exist_ok=True)
    WORK, OUT, USER = S.WORK, S.OUT, S.USER
    import jobs
    if args.list:
        for n in sorted(jobs.JOBS):
            print("%-10s %s" % (n, (jobs.JOBS[n].__doc__ or "").strip().splitlines()[0]))
        return 0
    if not os.path.exists("/tmp/.X11-unix/X" + DISPLAY.lstrip(":")):
        sys.exit("display %s not running, start it with start_display.sh" % DISPLAY)

    failed = []
    for n in args.jobs or sorted(jobs.JOBS):
        print("job", n)
        S.cleanup()
        try:
            jobs.JOBS[n]()
        except Exception as e:
            print("  FAILED:", e)
            failed.append(n)
        S.cleanup()

    if args.install:
        for f in sorted(glob.glob(os.path.join(OUT, "Figure[0-9][0-9].png"))):
            shutil.copy(f, os.path.join(BNC_ROOT, "src", "IMG"))
            print("  installed", os.path.basename(f))
    if failed:
        print("failed:", " ".join(failed))
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())