argparse_wrap.py

# ARGPARSE_WRAP: simple form of command registration for argparse

# Open sourced from Vast.ai internal tools
# MIT license

import argparse
import sys
import os


class argument(object):
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
REMAINDER=argparse.REMAINDER
class apwrap(object):
    REMAINDER=argparse.REMAINDER
    def __init__(self, *args, **kwargs):
        self.parser = argparse.ArgumentParser(*args, **kwargs)
        self.parser.set_defaults(func=self.fail_with_help)
        self.subparsers_ = None
        self.added_help_cmd = False
        self.post_setup = []

    def fail_with_help(self, *a, **kw):
        self.parser.print_help(sys.stderr)
        raise SystemExit

    def add_argument(self, *a, **kw):
        return self.parser.add_argument(*a, **kw)

    def subparsers(self, *a, **kw):
        if self.subparsers_ is None:
            kw["metavar"] = "command"
            kw["help"] = "command to run. one of:"
            self.subparsers_ = self.parser.add_subparsers(*a, **kw)
        return self.subparsers_

    def command(self, *arguments, aliases=(), help=None):
        help_ = help
        if not self.added_help_cmd:
            self.added_help_cmd = True

            @self.command(argument("subcommand", default=None, nargs="?"), help="print this help message")
            def help(*a, **kw):
                self.fail_with_help()

        def inner(func):
            name = func.__name__.replace("_", "-")
            sp = self.subparsers().add_parser(name, aliases=aliases, help=help_)
            for arg in arguments:
                sp.add_argument(*arg.args, **arg.kwargs)
            sp.set_defaults(func=func)
            return func
        if len(arguments) == 1 and type(arguments[0]) != argument:
            func = arguments[0]
            arguments = []
            return inner(func)
        return inner

    def parse_args(self, *a, **kw):
        args = self.parser.parse_args(*a, **kw)
        for func in self.post_setup:
            func(args)
        return args

requirements.txt

blessed
watchdog
psutil

watch.py

# process restarter, for development.
# especially relevant for deep process trees,
# see reap_children().

# Open sourced from Vast.ai internal tools
# MIT license

# examples:
# python watch.py watch 'src/**.jsx' 'src/**.scss' --command yarn build
# python watch.py autorestart 'src/**.jsx' 'src/**.scss' --command yarn run
# python watch.py watch 'src/**.cpp' 'src/**.h' --command make all
# python watch.py autorestart 'src/**.cpp' 'src/**.h' --command bash -c 'make all; ./target'

from argparse_wrap import apwrap, argument
import textwrap
import sys
import collections
import shlex
import subprocess
import signal
import os
import time
import blessed
from watchdog.observers import Observer
from watchdog.tricks import Trick
import psutil
import argparse
term = blessed.Terminal()
apw = apwrap()

os.environ.pop("MAKEFLAGS", None)
os.environ.pop("MAKELEVEL", None)
os.environ.pop("MAKE_TERMERR", None)
os.environ.pop("MAKE_TERMOUT", None)
os.environ.pop("MFLAGS", None)
os.environ["_"] = "watch"

quiet = False

from watchdog.events import FileSystemEvent
real_init = FileSystemEvent.__init__
def event_init(self, *a, **kw):
    real_init(self, *a, **kw)
    self.when = time.time()
FileSystemEvent.__init__ = event_init

def print_(*a, **kw):
    if quiet:
        return
    print(*a, **kw)


def handle_sigterm(_signum, _frame):
    raise KeyboardInterrupt()
signal.signal(signal.SIGTERM, handle_sigterm)


class CommandTrick(Trick):

    """Executes shell commands in response to matched events."""

    def __init__(self, cmd, patterns=None, ignore_patterns=None,
                 ignore_directories=False):
        super(CommandTrick, self).__init__(patterns, ignore_patterns,
                                                ignore_directories)
        self.process = None
        self.command = cmd
        self.on_any_event(None)

    def on_any_event(self, event):
        time.sleep(0.1)
        subprocess.call(self.command)

procname_cache = {}

class ProcName:
    def __init__(self, p):
        self.depth = 0
        if p is None:
            self.name = "[missing]"
            self.cmdline = ["[missing]"]
        else:
            try:
                self.name = p.name()
            except psutil.Error:
                self.name = "[error]"
            try:
                self.cmdline = p.cmdline()
            except psutil.Error:
                self.cmdline = ["[error]"]
        self.cmdline_fmt = " ".join(shlex.quote(x) for x in self.cmdline)
        self.trimmed = self.cmdline_fmt
        if len(self.trimmed) > 40:
            self.trimmed = self.trimmed[:36] + " ..."

procname_missing = ProcName(None)

def format_process(p, pid=None, indent=">", box=True):
    if p is None and pid is None:
        return "[<< totally missing wtf >>]"
    if pid is None:
        pid = p.pid
    procname = procname_cache.get(pid, procname_missing)
    if p is None:
        status = "missing"
    else:
        try:
            status = p.status()
        except psutil.NoSuchProcess:
            status = "gone"
        except psutil.ZombieProcess:
            status = "zombie"
        except psutil.Error:
            status = "error"
    indented = indent*procname.depth
    if box:
        prefix = "{}{}".format(indented, str(pid)).ljust(12)
        res = "[{}: {} ({})]".format(prefix, procname.trimmed, status)
        res = res.ljust(67)
        return term.blue+term.bold+res+term.bold+term.white
    else:
        res = "{}: {} ({})".format(str(pid), procname.cmdline_fmt, status)
        wrapped = textwrap.wrap(res, width=term.width-len(indented)*2)
        final = [indented + wrapped[0]]
        for line in wrapped[1:]:
            final.append(" "*(len(indented)+4)+"... "+line)
        return term.blue+term.bold+"\n".join(final)+term.bold+term.white

def print_tree(rootpid, tree, info, depth=0):
    proc = info.get(rootpid, None)
    procname = procname_cache.get(rootpid, None)
    if procname:
        procname.depth = depth

    print_(term.normal+format_process(proc, rootpid, indent="  > ", box=False)+term.normal)
    if rootpid not in tree:
        return

    for child in tree[rootpid]:
        print_tree(child, tree, info, depth+1)

def reap_children():
    print_(term.bold+term.white+"WATCH: killing process subtree..."+term.normal)
    term_printed = set()
    def on_terminate(proc):
        if proc in term_printed:
            return
        term_printed.add(proc)
        print_(term.bold+term.white+"WATCH: {}: terminated with exit code {}".format(format_process(proc), proc.returncode)+term.normal)

    me = psutil.Process()
    procs = me.children(recursive=True)
    procname_cache.clear()

    try:
        tree = collections.defaultdict(list)
        info = collections.defaultdict(lambda: None)
        info[me.pid] = me
        procname_cache[me.pid] = ProcName(me)
        nonroots = set()
        for p in procs:
            try:
                info[p.pid] = p
                nonroots.add(p.pid)
                tree[p.ppid()].append(p.pid)
                procname_cache[p.pid] = ProcName(p)
            except (psutil.NoSuchProcess, psutil.ZombieProcess):
                pass
        for root in (tree.keys() - nonroots):
            print_tree(root, tree, info)
    except:
        import traceback
        traceback.print_exc()
        print_(term.bold+term.white+'^^^^ error printing tree ^^^^'+term.normal)

    permission_denied = set()
    for p in procs:
        try:
            p.send_signal(signal.SIGINT)
            print_(term.bold+term.white+"WATCH: {}: sent sigint".format(format_process(p))+term.normal)
        except (OSError, psutil.NoSuchProcess):
            print_(term.bold+term.white+"WATCH: {}: not running by the time we tried to interrupt it".format(format_process(p))+term.normal)
        except psutil.AccessDenied:
            permission_denied.add(p)
            print_(term.bold+term.white+"WATCH: {}: can't kill, got a permission error. won't wait long for it.".format(format_process(p))+term.normal)
    print_(term.bold+term.white+"WATCH: Waiting for processes to die..."+term.normal)
    gone, alive = psutil.wait_procs(procs, timeout=1, callback=on_terminate)
    for p in permission_denied:
        print_(term.bold+term.white+"WATCH: {}: didn't die, not waiting any longer for it".format(format_process(p))+term.normal)
        procs.remove(p)
    gone, alive = psutil.wait_procs(procs, timeout=0.01, callback=on_terminate)
    if not alive:
        print_(term.bold+term.white+"WATCH: done killing subtree"+term.normal)
        return

    # send SIGTERM
    for p in alive:
        try:
            p.terminate()
            print_(term.bold+term.white+"WATCH: {}: sent sigterm".format(format_process(p))+term.normal)
        except (OSError, psutil.NoSuchProcess):
            print_(term.bold+term.white+"WATCH: {}: not running by the time we tried to terminate it".format(format_process(p))+term.normal)
    print_(term.bold+term.white+"WATCH: Waiting for processes to die..."+term.normal)
    gone, alive = psutil.wait_procs(alive, timeout=1, callback=on_terminate)
    if not alive:
        print_(term.bold+term.white+"WATCH: done killing subtree"+term.normal)
        return

    # send SIGKILL
    for p in alive:
        print_(term.bold+term.white+"WATCH: {}: survived SIGTERM; trying SIGKILL".format(format_process(p))+term.normal)
        try:
            p.kill()
        except (OSError, psutil.NoSuchProcess):
            print_(term.bold+term.white+"WATCH: {}: not running by the time we tried to kill it".format(format_process(p))+term.normal)
    print_(term.bold+term.white+"WATCH: Waiting for processes to die..."+term.normal)
    gone, alive = psutil.wait_procs(alive, timeout=4, callback=on_terminate)
    if not alive:
        print_(term.bold+term.white+"WATCH: done killing subtree"+term.normal)
        return

    # give up
    for p in alive:
        print_(term.bold+term.white+"WATCH: {}: survived SIGKILL; giving up".format(format_process(p))+term.normal)


class AutoRestartTrick(Trick):
    def __init__(self, command, patterns, ignore_patterns, ignore_directories):
        super(AutoRestartTrick, self).__init__(
            patterns, ignore_patterns, ignore_directories)
        self.command = command
        self.process = None
        self.ratelimit = 0.0
        self.last_event = time.time()
        self.wait_for_reemit = False
        self.unseen_events = 0

    def start(self):
        print_(term.bold+term.white+"WATCH: STARTING PROCESS..."+term.normal)
        if self.process is not None:
            print_("\033[31mWATCH: WARNING: TRIED TO START PROCESS TWICE\033[m")
        self.process = subprocess.Popen(self.command, preexec_fn=os.setsid)
        print_(term.bold+term.green+"WATCH: PROCESS SHOULD NOW BE RUNNING"+term.normal)

    def stop(self):
        print_(term.bold+term.yellow+"WATCH: STOPPING PROCESS..."+term.normal)
        if self.process is None:
            print_(term.bold+term.white+"WATCH: HAVE NOT LAUNCHED PROCESS"+term.normal)
            return
        reap_children()
        self.process = None
        print_(term.bold+term.white+"WATCH: MARKED AS NO PROCESS"+term.normal)

    def on_any_event(self, event):
        min_wait = 0.2
        thresh_wait = 0.02

        self.stop()

        now = event.when

        delta = max(now - self.last_event, 1e-3)
        self.last_event = now
        self.ratelimit = min(2, max(0, (0.96 ** delta)*self.ratelimit + 3*(1.0/delta + 0.125)))

        time_to_wait = max(min_wait, self.ratelimit - delta)

        if time_to_wait > thresh_wait:
            print_(term.bold+term.orange+"WATCH: got event after ({:0.3f}s out of {:0.3f}s), ignore duration {:0.3f}s".format(delta, self.ratelimit, time_to_wait)+term.normal)
            time.sleep(time_to_wait)
            print_(term.bold+term.orange+"WATCH: Skipping {} events".format(self.observer.event_queue.qsize())+term.normal)
            self.observer.event_queue.queue.clear()
        self.start()

@apw.command(
    argument('paths', nargs="+"),
    argument('-q', '--quiet', action="store_true"),
    argument('--command', nargs=argparse.REMAINDER, required=True),
)
def autorestart(args):
    global quiet
    quiet = args.quiet
    import signal

    patterns, ignores, dirnames = parse_patterns(args.paths)
    handler = AutoRestartTrick(command=args.command,
                               patterns=patterns,
                               ignore_patterns=ignores,
                               ignore_directories=True)
    handler.start()
    observer = Observer(timeout=1)
    run(handler, dirnames)
    handler.stop()

@apw.command(
    argument('paths', nargs="+"),
    argument('--command', nargs=argparse.REMAINDER, required=True),
)
def watch(args):
    patterns, ignores, dirnames = parse_patterns(args.paths)
    handler = CommandTrick(args.command, patterns=patterns,
                                ignore_patterns=ignores, #ignore_patterns,
                                ignore_directories=True)
    run(handler, dirnames)

def parse_patterns(paths):
    patterns = []
    dirnames = {}
    ignores = []
    for pathpattern in paths:
        if pathpattern.startswith("-"):
            pathpattern = pathpattern[1:]
            assert "/" not in pathpattern
            ignores.append(pathpattern)
        elif "**" in pathpattern:
            before, stars, after = pathpattern.partition("**")
            dirname = os.path.abspath(os.path.expanduser(os.path.dirname(before)))
            base = os.path.basename(before)
            dirnames[dirname] = True
            patterns.append(base+stars+after)
        else:
            dirname = os.path.dirname(pathpattern)
            basename = os.path.basename(pathpattern)
            assert "*" not in dirname
            dirnames.setdefault(dirname, False)
            patterns.append(basename)
    return patterns, ignores, dirnames

def run(handler, dirnames):
    observer = Observer(timeout=1.0)
    handler.observer = observer
    watches = []
    handler.watches = watches
    for dirname, recursive in dirnames.items():
        watches.append(observer.schedule(handler, os.path.abspath(os.path.expanduser(dirname)), recursive))
    try:
        observer.start()
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            print_("Stopping")
            observer.stop()
            print_("Done stopping")
        observer.join()
    except KeyboardInterrupt:
        print_("Shutting down forcefully?")
        pass
        print_("Stuff may or may not actually happen")


def main():
    args = apw.parse_args()
    args.func(args)

if __name__ == "__main__":
    main()