jw-pkg/src/python/jw/pkg/lib/App.py

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)