from __future__ import annotations import asyncio import cProfile import os import sys from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace from typing import TYPE_CHECKING, Any, cast, override from .AsyncRunner import AsyncRunner from .log import ( DEBUG, ERR, NOTICE, LogFlag, log, log_m, parse_log_flags, parse_log_level, set_log_flags, set_log_level ) from .Types import LoadTypes from .util import pretty_cmd if TYPE_CHECKING: import types from collections.abc import Awaitable, Collection from typing import TypeVar T = TypeVar('T') class App: # export def _add_arguments(self, parser: ArgumentParser) -> None: self.__parser.add_argument( '--log-flags', help = 'Log flags', default = self.__default_log_flags ) self.__parser.add_argument( '--log-level', help = 'Log level', default = self.__default_log_level ) self.__parser.add_argument( '--log-file', help = 'Log file', default = self.__default_log_file ) self.__parser.add_argument( '--backtrace', help = 'Show exception backtraces', action = 'store_true', default = self.__back_trace, ) self.__parser.add_argument( '--write-profile', help = 'Profile code and store output to file', default = None, ) def _default_log_flags_env(self) -> str: return 'JW_DEFAULT_LOG_FLAGS' def _default_log_flags(self, default: LogFlag) -> LogFlag: return default def _default_log_level_env(self) -> str: return 'JW_DEFAULT_LOG_LEVEL' def _default_log_level(self, default: int) -> int: return default def _default_log_file_env(self) -> str: return 'JW_DEFAULT_LOG_FILE' def _default_log_file(self, default: str | None) -> str | None: return default def _default_show_backtrace_env(self) -> str: return 'JW_DEFAULT_SHOW_BACKTRACE' def _default_show_backtrace(self, default: bool) -> bool: return default def __init__( self, description: str = '', name_filter: str = '^Cmd.*', modules: list[str] | None = None, eloop: None = None, ) -> None: def add_cmd_to_parser(cmd: AbstractCmd, parsers: Any) -> ArgumentParser: parser = cast( 'ArgumentParser', parsers.add_parser( cmd.name, help = cmd.help, description = cmd.description, aliases = cmd.aliases, formatter_class = ArgumentDefaultsHelpFormatter, ) ) parser.set_defaults(func = cmd.run) cmd.add_arguments(parser) cmd.set_parser(parser) return parser def add_cmds_to_parser( parent: AbstractCmd | App, parser: ArgumentParser, cmds: Collection[AbstractCmd], all: bool = False ) -> None: if not cmds: return class SubCommand: def __init__(self, cmd: AbstractCmd, parser: Any): self.cmd = cmd self.parser = parser title = 'Available subcommands' if isinstance(parent, AbstractCmd): title += ' of ' + parent.name subparsers = parser.add_subparsers( title = title, metavar = '', dest = 'command' ) scs: dict[str, SubCommand] = {} for cmd in cmds: cmd.set_parent(parent) scs[cmd.name] = SubCommand(cmd, add_cmd_to_parser(cmd, subparsers)) for alias in cmd.aliases: scs[alias] = scs[cmd.name] if all: seen: set[int] = set() for sc in scs.values(): if id(sc) not in seen: seen.add(id(sc)) add_cmds_to_parser( sc.cmd, sc.parser, sc.cmd.children, all = all ) return args, _ = self.__parser.parse_known_args() cmd_name = getattr(args, 'command', None) if cmd_name in scs: sc = scs[cmd_name] add_cmds_to_parser(sc.cmd, sc.parser, sc.cmd.children, all = all) from .Cmd import AbstractCmd self.__args: Namespace | None = None self.__cmdline: str | None = None self.__default_log_flags = self._default_log_flags( LogFlag.STDERR | LogFlag.POSITION | LogFlag.PRIO | LogFlag.COLOR ) if (env := os.getenv(self._default_log_flags_env(), None)) is not None: self.__default_log_flags = parse_log_flags(env) self.__default_log_level = self._default_log_level(NOTICE) if (env := os.getenv(self._default_log_level_env(), None)) is not None: self.__default_log_level = parse_log_level(env) self.__default_log_file = self._default_log_file(None) if (env := os.getenv(self._default_log_file_env(), None)) is not None: self.__default_log_file = env self.__back_trace = self._default_show_backtrace(False) if (env := os.getenv(self._default_show_backtrace_env())) is not None: self.__back_trace = env.lower() in ['1', 'true'] set_log_flags(self.__default_log_flags) set_log_level(self.__default_log_level) self.__eloop = eloop self.__own_eloop = False if eloop is None: self.__eloop = asyncio.get_event_loop() self.__own_eloop = True self.__async_runner: AsyncRunner | None = None self.__parser = ArgumentParser( formatter_class = ArgumentDefaultsHelpFormatter, description = description, add_help = False, ) self._add_arguments(self.__parser) args, _ = self.__parser.parse_known_args() set_log_flags(args.log_flags) set_log_level(args.log_level) log(DEBUG, f'-------------- Running: >{pretty_cmd(sys.argv)}<') cmd_classes: LoadTypes[AbstractCmd] = LoadTypes( modules if modules else ['__main__'], type_name_filter = name_filter, type_filter = [AbstractCmd], ) add_all_parsers = ( '-h' in sys.argv or '--help' in sys.argv or '_ARGCOMPLETE' in os.environ ) add_cmds_to_parser( self, self.__parser, [cmd_class(self) for cmd_class in cmd_classes], all = add_all_parsers, ) # -- Add help only now, wouldn't want to have parse_known_args() exit # on --help with subcommands missing self.__parser.add_argument( '-h', '--help', action = 'help', help = 'Show this help message and exit' ) def __del__(self) -> None: if self.__own_eloop: if self.__eloop is not None: self.__eloop.close() self.__eloop = None self.__own_eloop = False async def __aenter__(self) -> None: pass async def __aexit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: types.TracebackType | None, ) -> None: pass async def __run(self, argv: list[str] | None = None) -> None: try: # Import argcomplete only here to not require it to be compatible # with minimal environments from argcomplete.completers import BaseCompleter # type: ignore[import-not-found, unused-ignore] # isort: skip class NoopCompleter(BaseCompleter): # type: ignore[misc, unused-ignore] @override def __call__( # pyright: ignore[reportGeneralTypeIssues] self, *args: Any, **kwargs: Any) -> None: return None import argcomplete # type: ignore[import-not-found, unused-ignore] argcomplete.autocomplete(self.__parser, default_completer = NoopCompleter()) except Exception: pass self.__args = self.__parser.parse_args(args = argv) set_log_flags(self.__args.log_flags) set_log_level(self.__args.log_level) self.__back_trace = self.__args.backtrace exit_status = 0 pr = None if self.__args.write_profile is None else cProfile.Profile() if pr is not None: pr.enable() try: ret = await self._run(self.__args) if isinstance(ret, int) and ret >= 0 and ret <= 0xFF: exit_status = ret except Exception as e: log_m(ERR, f'Failed: {repr(e) if self.__back_trace else str(e)}') exit_status = 1 # AssertionErrors are programming errors, hence a programmer should # get a chance to figure it out if self.__back_trace or isinstance(e, AssertionError): raise finally: if pr is not None: pr.disable() log( NOTICE, f'Writing profile statistics to {self.__args.write_profile}' ) assert self.__args.write_profile is not None, 'args.write_profile' pr.dump_stats(self.__args.write_profile) if exit_status: sys.exit(exit_status) # Do the main work. Tries to run sub-commands by default. Overwrite if you # want to do something else, for instance if you don't have sub-commands, # or if want to do anything before and / or after the subcommands. async def _run(self, args: Namespace) -> None | int: if not hasattr(self.__args, 'func'): self.__parser.print_help() return None # Run sub-command. Overwrite if you want to do anything before or after return cast('None | int', await self.args.func(args)) def call_async(self, awaitable: Awaitable[T], timeout: float | None = None) -> T: return self.async_runner.call(awaitable, timeout) @property def eloop(self) -> asyncio.AbstractEventLoop: if self.__eloop is None: raise Exception('Tried to get inexistent event loop from application') return self.__eloop @property def async_runner(self) -> AsyncRunner: if self.__async_runner is None: self.__async_runner = AsyncRunner() return self.__async_runner @property def cmdline(self) -> str: if self.__cmdline is None: import shlex with open('/proc/self/cmdline', 'rb') as f: raw = f.read().split(b'\0')[:-1] self.__cmdline = ' '.join(shlex.quote(arg.decode()) for arg in raw) return self.__cmdline @property def args(self) -> Namespace: if self.__args is None: raise Exception('Tried to get inexistent argument list from application') return self.__args @property def parser(self) -> ArgumentParser: return self.__parser def run(self, argv: list[str] | None = None) -> None: try: ret = self.eloop.run_until_complete(self.__run(argv)) finally: if self.__async_runner: self.__async_runner.close() self.__async_runner = None return ret def run_sub_commands( # export description: str = '', name_filter: str = '^Cmd.*', modules: list[str] | None = None, argv: list[str] | None = None, ) -> None: app = App(description, name_filter, modules) app.run(argv = argv)