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

from __future__ import annotations

import asyncio
import cProfile
import os
import sys
import warnings

from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace
from typing import TYPE_CHECKING, Any, NamedTuple, cast, override

from .AsyncRunner import AsyncRunner
from .Cmd import AbstractCmd
from .log import (
    DEBUG,
    ERR,
    NOTICE,
    WARNING,
    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')

def _get_current_event_loop() -> asyncio.AbstractEventLoop | None:
    """Return the current event loop of this thread, or None if there is
    none, without creating one implicitly or emitting a deprecation
    warning."""
    with warnings.catch_warnings():
        warnings.simplefilter('error', DeprecationWarning)
        try:
            return asyncio.get_event_loop()
        except (RuntimeError, DeprecationWarning):
            return None

class _SubCommand(NamedTuple):

    cmd: AbstractCmd
    parser: ArgumentParser

class App:  # export

    def _add_arguments(self, parser: ArgumentParser) -> None:
        parser.add_argument(
            '--log-flags',
            help = 'Log flags',
            default = self.__default_log_flags,
            type = parse_log_flags,
        )
        parser.add_argument(
            '--log-level',
            help = 'Log level',
            default = self.__default_log_level,
            type = parse_log_level,
        )
        parser.add_argument(
            '--log-file', help = 'Log file', default = self.__default_log_file
        )
        parser.add_argument(
            '--backtrace',
            help = 'Show exception backtraces',
            action = 'store_true',
            default = self.__back_trace,
        )
        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: asyncio.AbstractEventLoop | None = None,
    ) -> None:

        self.__args: Namespace | None = None
        self.__cmdline: str | None = None
        self.__description = description

        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.__async_runner: AsyncRunner | None = None
        self.__eloop = eloop
        self.__own_eloop = False

        cmd_classes: LoadTypes[AbstractCmd] = LoadTypes(
            modules if modules else ['__main__'],
            type_name_filter = name_filter,
            type_filter = [AbstractCmd],
        )
        self.__cmds: list[AbstractCmd] = [cmd_class(self) for cmd_class in cmd_classes]
        self._build_parser()

    def _build_parser(self, argv: list[str] | 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

            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)
                if cmd.name in scs:
                    log(WARNING, f'Duplicate subcommand name: {cmd.name}')
                scs[cmd.name] = _SubCommand(cmd, add_cmd_to_parser(cmd, subparsers))
                for alias in cmd.aliases:
                    if alias != cmd.name and alias in scs:
                        log(
                            WARNING,
                            f'Subcommand alias "{alias}" of "{cmd.name}" '
                            'collides with an earlier subcommand',
                        )
                    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
            # -- Re-parse the command line to find the invoked subcommand.
            # This works because every level below uses dest = 'command',
            # so each pass descends one level further into the command
            # tree.
            args, _ = self.__parser.parse_known_args(argv)
            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)

        cmdline = sys.argv if argv is None else argv
        if argv is None:
            argv = sys.argv[1:]
        add_all_parsers = (
            '-h' in argv or '--help' in argv or '_ARGCOMPLETE' in os.environ
        )

        self.__parser = ArgumentParser(
            formatter_class = ArgumentDefaultsHelpFormatter,
            description = self.__description,
            add_help = False,
        )
        self._add_arguments(self.__parser)

        args, _ = self.__parser.parse_known_args(argv)
        set_log_flags(args.log_flags)
        set_log_level(args.log_level)

        log(DEBUG, f'-------------- Running: >{pretty_cmd(cmdline)}<')

        add_cmds_to_parser(self, self.__parser, self.__cmds, 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 close(self) -> None:
        """Close the application and release all resources"""
        if self.__async_runner is not None:
            self.__async_runner.close()
            self.__async_runner = None
        if self.__own_eloop:
            if self.__eloop is not None:
                if not self.__eloop.is_closed():
                    self.__eloop.close()
                self.__eloop = None
            self.__own_eloop = False

    async def __aenter__(self) -> App:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: types.TracebackType | None,
    ) -> None:
        self.close()

    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 ImportError:
            log(DEBUG, 'argcomplete is not installed, shell completion disabled')
        except Exception as e:
            log(DEBUG, f'Shell completion disabled: {e}')

        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:
            result = await self._run(self.__args)
            if isinstance(result, int):
                if 0 <= result <= 0xFF:
                    exit_status = result
                else:
                    log(
                        WARNING,
                        f'Command returned invalid exit status {result}, '
                        'using 1 instead',
                    )
                    exit_status = 1
        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(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 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:
        previous_eloop: asyncio.AbstractEventLoop | None = None
        if self.__eloop is None:
            previous_eloop = _get_current_event_loop()
            eloop = asyncio.new_event_loop()
            asyncio.set_event_loop(eloop)
            self.__eloop = eloop
            self.__own_eloop = True
        try:
            if argv is not None:
                self._build_parser(argv)
            ret = self.eloop.run_until_complete(self.__run(argv))
        finally:
            self.close()
            # -- Restore the event loop the thread had before run(), or
            # unset the loop if there was none.
            if previous_eloop is not None:
                asyncio.set_event_loop(previous_eloop)
            else:
                asyncio.set_event_loop(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)