SERELIB/1
name=argparser
version=0.1.1
entry=src/lib.sere

FILE src/lib.sere 6703

import os

class ArgumentError(Exception):
    pass

enum ArgumentAction:
    Store
    StoreTrue

def clean_name(name: str) -> str:
    if name.starts_with("--"):
        return name[2:]
    if name.starts_with("-"):
        return name[1:]
    return name

class Argument:
    @frozen name: str
    @frozen aliases: list[str]
    @frozen help: str
    @frozen required: bool
    @frozen default: str
    @frozen action: ArgumentAction

    def __init__(
        self,
        name: str,
        aliases: list[str] | None = None,
        help: str = "",
        required: bool = False,
        default: str = "",
        action: ArgumentAction = ArgumentAction.Store,
    ) -> None:
        self.name = name
        self.aliases = do:
            if aliases is None:
                list[str]()
            else:
                aliases
        self.help = help
        self.required = required
        self.default = default
        self.action = action

class Namespace:
    @private values: dict[str, str]
    @private flags: dict[str, bool]

    def __init__(self) -> None:
        self.values = {}
        self.flags = {}

    def set(self, name: str, value: str) -> None:
        self.values[name] = value

    def set_bool(self, name: str, value: bool) -> None:
        self.flags[name] = value

    def get(self, name: str) -> str:
        if name in self.values:
            return self.values[name]
        return ""

    def get_bool(self, name: str) -> bool:
        if name in self.flags:
            return self.flags[name]
        return False

    def has(self, name: str) -> bool:
        if (name in self.values) or (name in self.flags):
            return True
        return False

class ArgParser:
    @frozen program: str
    @frozen description: str
    arguments: list[Argument]

    def __init__(self, program: str, description: str = "") -> None:
        self.program = program
        self.description = description
        self.arguments = []

    def add_argument(
        self,
        name: str,
        aliases: list[str] | None = None,
        help: str = "",
        required: bool = False,
        default: str = ""
    ) -> None:
        argument = Argument(
            name,
            aliases=aliases,
            help=help,
            required=required,
            default=default,
            action=ArgumentAction.Store
        )

        self.arguments.append(argument)

    def add_flag(
        self,
        name: str,
        aliases: list[str] | None = None,
        help: str = "",
    ) -> None:
        argument = Argument(
            name,
            aliases=aliases,
            help=help,
            required=False,
            default="",
            action=ArgumentAction.StoreTrue
        )

        self.arguments.append(argument)

    def find_argument(self, token: str) -> Argument | None:
        for argument in self.arguments:
            if token == argument.name:
                return argument

            for alias in argument.aliases:
                if alias == token:
                    return argument
        return None

    def create_namespace(self) -> Namespace:
        result = Namespace()
        for argument in self.arguments:
            key = clean_name(argument.name)

            match argument.action:
                case ArgumentAction.Store:
                    if argument.default != "":
                        result.set(key, argument.default)
                case ArgumentAction.StoreTrue:
                    result.set_bool(key, False)
        return result

    def parse_args(self, argv: list[str]) -> Namespace:
        result = self.create_namespace()
        index: i32 = 1
        while index < len(argv):
            token = argv[index]

            argument = do:
                arg = self.find_argument(token)
                if arg is None:
                    raise ArgumentError("unknown argument: " + token)
                arg as Argument

            key = clean_name(argument.name)

            match argument.action:
                case ArgumentAction.StoreTrue:
                    result.set_bool(key, True)
                case ArgumentAction.Store:
                    if index + 1 >= len(argv):
                        raise ArgumentError("expected value after " + token)

                    value = argv[index + 1]
                    result.set(key, value)
                    index += 1
            index += 1
        self.validate_required(result)
        return result

    def validate_required(self, namespace: Namespace) -> None:
        for argument in self.arguments:
            if not argument.required:
                continue

            key = clean_name(argument.name)

            if not namespace.has(key):
                raise ArgumentError(
                    "required argument missing: " + argument.name
                )

    def argument_names(self, argument: Argument) -> str:
        names = argument.name
        for alias in argument.aliases:
            names = alias + ", " + names
        return names

    def argument_summary(self, argument: Argument) -> str:
        if argument.action == ArgumentAction.StoreTrue:
            return "flag"
        if argument.required:
            return "required"
        if argument.default != "":
            return f"default: {argument.default}"
        return "optional"

    def print_help(self) -> void:
        print(f"usage: {self.program} [options]")
        if self.description != "":
            print("")
            print(self.description)

        print("")
        print("options:")
        print("\t-h, --help")
        print("\t\tShow this help message")

        for argument in self.arguments:
            names = self.argument_names(argument)
            print(f"\t{names}")

            label = self.argument_summary(argument)
            if argument.help != "":
                print(f"\t\t{argument.help} ({label})")
            else:
                print(f"\t\t{label}")

    def print_usage(self) -> void:
        print(f"usage: {self.program} [options]")

    def error(self, message: str) -> None:
        print(f"{self.program}: error: {message}")
        self.print_usage()
        print(f"run '{self.program} --help' for more information")

    def parse_or_exit(self, argv: list[str]) -> Namespace:
        for token in argv:
            if token == "-h" or token == "--help":
                self.print_help()
                os.exit(0)

        namespace: Namespace = Namespace()
        try:
            namespace = self.parse_args(argv)
        except ArgumentError as error:
            self.error(error.message)
            os.exit(2)
        return namespace