argparser.py

""" argparser - circlebot's non-failing argument parser meant to replace the `argparse` standard library. """
from __future__ import annotations

import re
import shlex
from dataclasses import dataclass
from enum import StrEnum
from typing import Any
from typing import Callable
from typing import Dict
from typing import List
from typing import Tuple
from typing import TypeVar

URI_REGEX = r"https?:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()@:%_\+.~#?&//=]*)"


class _MISSING:
    pass


MISSING = _MISSING()
T = TypeVar("T")


class ArgumentError(Exception): ...


class ArgumentActionType(StrEnum):
    STORE = "store"
    STORE_TRUE = "store_true"


@dataclass
class ArgumentType:
    names_or_flags: List[str]
    store_type: T
    action: ArgumentActionType
    _default: T | _MISSING = MISSING
    converter: Callable[[str], T] | None = None
    regex: str | None = None
    ignore_default: T | None = None
    is_parsed: bool = False
    nullable: bool = False

    @property
    def default(self) -> T | None:
        if self._default is MISSING:
            return None

        return self._default

    @property
    def name(self) -> str:
        return max(self.names_or_flags, key=len).strip("-")


class ArgumentParser:
    def __init__(self, *, skip_uris: bool = False) -> None:
        self.skip_uris: bool = skip_uris
        self.keys: List[ArgumentType] = []

    def add_argument(
        self,
        *names_or_flags: List[str],
        action: ArgumentActionType = ArgumentActionType.STORE,
        _type: type = str,
        default: T | _MISSING = MISSING,
        converter: Callable[[str], T] | None = None,
        regex: str | None = None,
        ignore_default: T | None = None,
        nullable: bool = False,
    ) -> None:
        if action not in (ArgumentActionType.STORE, ArgumentActionType.STORE_TRUE):
            raise ArgumentError(f"Unknown action {action}")

        if len(names_or_flags) == 0:
            raise ArgumentError("Can't parse argument with 0 names (or flags).")

        if action == ArgumentActionType.STORE_TRUE:
            _type = bool

        self.keys.append(
            ArgumentType(
                names_or_flags,
                _type,
                action,
                default,
                converter,
                regex,
                ignore_default,
                nullable=nullable,
            ),
        )

    def parse(self, arguments: str) -> Tuple[Dict[str, Any], List[str]]:
        to_parse = shlex.split(arguments)

        parsed: Dict[str, Any] = {}
        unknown: List[str] = []
        found: List[str] = []
        # Index of the arg to skip.
        look_ahead = 0

        for idx, arg in enumerate(to_parse):
            # Skip args we manually skipped
            if look_ahead > idx:
                continue

            if self.skip_uris:
                if re.match(URI_REGEX, arg) is not None:
                    continue

            for key in self.keys:
                if key.is_parsed:
                    continue

                arg_lower = arg.lower()
                has_any_flags = any(fl.startswith("-") for fl in key.names_or_flags)
                if (
                    not any(arg_lower.startswith(flag) for flag in key.names_or_flags)
                    and has_any_flags
                ):
                    continue

                has_oneline_arg = any(
                    (arg_lower.startswith(flag) and arg_lower != flag)
                    for flag in key.names_or_flags
                )
                if has_any_flags and has_oneline_arg:
                    if (
                        key.action == ArgumentActionType.STORE_TRUE
                    ):  # we might've gotten the wrong key.
                        continue

                if key.action == ArgumentActionType.STORE_TRUE:
                    found.extend(key.names_or_flags)
                    parsed[key.name] = True
                    break

                if len(to_parse) >= idx + 1 or has_oneline_arg:
                    if has_oneline_arg or not has_any_flags:
                        arg_convert = arg
                        for flag in key.names_or_flags:
                            arg_convert = arg_convert.strip(flag)

                        if arg_convert.startswith("="):
                            arg_convert = arg_convert[1:]
                    else:
                        look_ahead = idx + 1

                        if look_ahead >= len(to_parse):
                            continue

                        arg_convert = to_parse[look_ahead]
                    arg_p: T | None = None

                    try:
                        if key.converter is not None:
                            arg_p = key.converter(arg_convert)

                            if (not arg_p and not key.nullable) or arg_p == key.ignore_default:
                                continue
                        elif key.regex is not None:
                            arg_p = re.match(key.regex, arg_convert)
                        else:
                            arg_p = key.store_type(arg_convert)

                        if key.store_type != type(arg_p) and arg_p is not None:
                            arg_p = key.store_type(arg_p)
                    except Exception:  # ignore any (and all) exceptions
                        continue

                    if arg_p is None:
                        continue

                    if key.store_type is str:
                        # Strip any unwanted quotes
                        arg_p = arg_p.strip('"').strip("'")

                    key.is_parsed = True
                    parsed[key.name] = arg_p
                    found.extend(key.names_or_flags)
                    found.append(arg_convert)
                    if has_oneline_arg:
                        found.append(arg)

            if arg not in found:
                unknown.append(arg)

        # append args that haven't been found
        for key in self.keys:
            if not any(flag in found for flag in key.names_or_flags):
                parsed[key.name] = key.default

        return parsed, unknown