scratch.py

from dataclasses import dataclass, field, fields
from copy import deepcopy
from typing import (
    reveal_type,
    overload,
    TypeVar,
    Any,
    Literal,
    TYPE_CHECKING,
    get_args,
    get_origin,
    Union,
    ClassVar,
)
from types import UnionType
from collections.abc import Mapping, Callable
from argparse import ArgumentParser
from enum import Enum, auto
from pathlib import Path


if TYPE_CHECKING:
    from _typeshed import DataclassInstance

_T = TypeVar("_T")

ActionType = Literal["store", "store_const", "append", "append_const", "count"] | None
NargsType = Literal["?", "+", "*"] | int | None


class _MISSING_TYPE:
    pass


MISSING = _MISSING_TYPE()

ARGPARSE_KEY = object()


class AddMemberName(Enum):
    NO = auto()
    YES = auto()
    HYPHEN = auto()


def argparse_field(
    # default is something that occurs both in dataclasses and argparse, the handling is a bit nuanced
    argparse_data: dict[str, Any],
    default: _T | _MISSING_TYPE = MISSING,
    # dataclass arguments
    init: bool = True,
    repr: bool = True,
    hash: bool | None = None,
    compare: bool = True,
    metadata: Mapping[Any, Any] | None = None,
    kw_only: bool = False,
) -> _T:
    if default is MISSING:
        default_dict: dict[str, Any] = {}
    else:
        default_dict = {"default_factory": lambda: deepcopy(default)}

    if metadata is None:
        metadata = dict()
    else:
        metadata = {k: v for k, v in metadata.items()}

    metadata[ARGPARSE_KEY] = argparse_data

    return field(
        **default_dict,
        init=init,
        repr=repr,
        hash=hash,
        compare=compare,
        metadata=metadata,
        kw_only=kw_only,
    )


@dataclass(init=True, eq=False)
class Exclusive:
    required: bool


def exclusive_group(required: bool = False) -> Exclusive:
    return Exclusive(required=required)


def subparsers(
    title: str | None = None,
    description: str | None = None,
    prog: str | None = None,
    parser_class: type | None = None,
    action: Any = None,
    *,
    default: _T | _MISSING_TYPE = MISSING,
    # dataclass arguments
    init: bool = True,
    repr: bool = True,
    hash: bool | None = None,
    compare: bool = True,
    metadata: Mapping[Any, Any] | None = None,
    kw_only: bool = False,
) -> Any:
    return argparse_field(
        {
            "is_subparser": True,
        },
        default=default,
        init=init,
        repr=repr,
        hash=hash,
        compare=compare,
        metadata=metadata,
        kw_only=kw_only,
    )


# If required is True, do not accept default argument at all
@overload
def option(
    short_arg: str | None = ...,
    *,
    more_arg_names: list[str] | None = ...,
    use_field_name: bool = True,
    exclusive_group: Exclusive | None = None,
    # dataclass arguments
    init: bool = ...,
    repr: bool = ...,
    hash: bool | None = ...,
    compare: bool = ...,
    metadata: Mapping[Any, Any] | None = ...,
    kw_only: bool = ...,
    # argparse arguments
    action: ActionType = ...,
    nargs: NargsType = ...,
    const: Any = ...,
    # skip default as we already handle that via dataclass
    type: Callable[[str], _T] | None = ...,
    choices: list[_T] | None = ...,
    required: Literal[True],
    help: str | None = ...,
    metavar: str | None = ...,
    deprecated: bool = ...,
) -> _T: ...


# If not required and no default provided, the default value is None
# return type is None (mypy can't seem to handle this without an overload)
@overload
def option(
    short_arg: str | None = ...,
    *,
    more_arg_names: list[str] | None = ...,
    use_field_name: bool = True,
    exclusive_group: Exclusive | None = ...,
    # dataclass arguments
    init: bool = ...,
    repr: bool = ...,
    hash: bool | None = ...,
    compare: bool = ...,
    metadata: Mapping[Any, Any] | None = ...,
    kw_only: bool = ...,
    # argparse arguments
    action: ActionType = ...,
    nargs: NargsType = ...,
    const: Any = ...,
    # skip default as we already handle that via dataclass
    type: Callable[[str], _T] | None = ...,
    choices: list[_T] | None = ...,
    required: Literal[False] = ...,
    help: str | None = ...,
    metavar: str | None = ...,
    deprecated: bool = ...,
) -> _T | None: ...

# Finally, the case where default is provided, and required is False
@overload
def option(
    short_arg: str | None = ...,
    *,
    more_arg_names: list[str] | None = ...,
    use_field_name: bool = True,
    exclusive_group: Exclusive | None = None,
    # default is something that occurs both in dataclasses and argparse, the handling is a bit nuanced
    default: _T | None = ...,
    # dataclass arguments
    init: bool = ...,
    repr: bool = ...,
    hash: bool | None = ...,
    compare: bool = ...,
    metadata: Mapping[Any, Any] | None = ...,
    kw_only: bool = ...,
    # argparse arguments
    action: ActionType = ...,
    nargs: NargsType = ...,
    const: Any = ...,
    # skip default as we already handle that via dataclass
    type: Callable[[str], _T] | None = ...,
    choices: list[_T] | None = ...,
    required: Literal[False] = ...,
    help: str | None = ...,
    metavar: str | None = ...,
    deprecated: bool = ...,
) -> _T: ...


def option(
    short_arg=None,
    *,
    more_arg_names=None,
    use_field_name=True,
    exclusive_group=None,
    # default is something that occurs both in dataclasses and argparse, the handling is a bit nuanced
    default=None,
    # dataclass arguments
    init: bool = True,
    repr: bool = True,
    hash=None,
    compare=True,
    metadata=None,
    kw_only: bool = False,
    # argparse arguments
    action: ActionType = None,
    nargs: NargsType = None,
    const: Any = None,
    # skip default as we already handle that via dataclass
    type=None,
    choices=None,
    required=False,
    help=None,
    metavar=None,
    deprecated=False,
):
    if use_field_name:
        add_member_name = AddMemberName.HYPHEN
    else:
        add_member_name = AddMemberName.NO

    if short_arg is None:
        short_list = []
    else:
        short_list = [short_arg]

    return argparse_field(
        {
            "default": default,
            "name_or_flags": short_list + (more_arg_names or []),
            "exclusive_group": exclusive_group,
            "add_member_name": add_member_name,
            "action": action,
            "nargs": nargs,
            "const": const,
            "type": type,
            "choices": choices,
            "required": required,
            "help": help,
            "metavar": metavar,
            "deprecated": deprecated,
        },
        default=default,
        init=init,
        repr=repr,
        hash=hash,
        compare=compare,
        metadata=metadata,
        kw_only=kw_only,
    )


def positional(
    # dataclass arguments
    *,
    default: _T | _MISSING_TYPE = MISSING,
    init: bool = True,
    repr: bool = True,
    hash: bool | None = None,
    compare: bool = True,
    metadata: Mapping[Any, Any] | None = None,
    kw_only: bool = False,
    # argparse arguments
    nargs: NargsType = None,
    # skip default as we already handle that via dataclass
    type: Callable[[str], _T] | None = None,
    choices: list[_T] | None = None,
    help: str | None = None,
    metavar: str | None = None,
    deprecated: bool = False,
) -> _T:
    return argparse_field(
        {
            "default": default,
            "add_member_name": AddMemberName.YES,
            "nargs": nargs,
            "type": type,
            "choices": choices,
            "help": help,
            "metavar": metavar,
            "deprecated": deprecated,
        },
        default=default,
        init=init,
        repr=repr,
        hash=hash,
        compare=compare,
        metadata=metadata,
        kw_only=kw_only,
    )


_D = TypeVar("_D", bound="DataclassInstance")

AUTO_CONVERTED_TYPES = {int, float, Path}


def add_args(
    d: type[_D], parser: ArgumentParser, subdest: str = ""
) -> tuple[dict[str, type], str]:
    exclusive_groups = {}
    subparser_field_name = ""
    nested_info = dict()
    for f in fields(d):
        argparse_data = dict(f.metadata[ARGPARSE_KEY])

        if argparse_data.pop("is_subparser", False):
            if subparser_field_name != "":
                raise ValueError(
                    "Only one subparser field is allowed in a dataclass. "
                    f"Found multiple: {subparser_field_name} and {f.name}"
                )
            subparser_field_name = f.name
            subparsers = parser.add_subparsers(dest=".subparser_name", **argparse_data)
            assert get_origin(f.type) in {UnionType, Union}
            for t in get_args(f.type):
                if t is type(None):
                    continue
                nested_info[t.subparser_name] = t
                subparser = subparsers.add_parser(t.subparser_name)
                add_args(t, subparser, ".subparsers.")

            continue

        add_member_name = argparse_data.pop("add_member_name")
        name_or_flags = argparse_data.pop("name_or_flags", [])
        # if add_member_name == AddMemberName.YES:
        #     name_or_flags = [f.name] + list(name_or_flags)
        if add_member_name == AddMemberName.HYPHEN:
            name_or_flags = ["--" + f.name.replace("_", "-")] + list(name_or_flags)

        argparse_data["dest"] = subdest + f.name

        typ = argparse_data.pop("type")
        if typ is None and f.type in AUTO_CONVERTED_TYPES:
            typ = f.type

        # A bit of a special case - if the field is in the form x: T | None, and the default is None,
        # then we still automatically provide the type argument, if T is suitable
        if typ is None and get_origin(f.type) in {UnionType, Union}:
            args = get_args(f.type)
            if (
                len(args) == 2
                and args[1] is type(None)
                and args[0] in AUTO_CONVERTED_TYPES
            ):
                typ = args[0]

        # required not being present either way means it's a positional argument
        # If nargs is not supplied, deduce whether or not to use ? based on whether
        # there is a default
        if (
            "required" not in argparse_data
            and argparse_data["nargs"] is None
            and argparse_data["default"] is not MISSING
        ):
            argparse_data["nargs"] = "?"

        if (eg_classvar := argparse_data.pop("exclusive_group", None)) is not None:
            if eg_classvar not in exclusive_groups:
                exclusive_groups[eg_classvar] = parser.add_mutually_exclusive_group(
                    required=eg_classvar.required
                )
            exclusive_groups[eg_classvar].add_argument(
                *name_or_flags, **argparse_data, type=typ
            )
        else:
            parser.add_argument(*name_or_flags, **argparse_data, type=typ)

    return nested_info, subparser_field_name


def parse_args(d: type[_D]) -> _D:
    parser = ArgumentParser()
    nested_info, subparser_field_name = add_args(d, parser)
    arg_dict = vars(parser.parse_args())
    if (nested_name := arg_dict.pop(".subparser_name", None)) is not None:
        nested_type = nested_info[nested_name]
        nested_dict = dict()
        for k, v in list(arg_dict.items()):
            if not k.startswith(".subparsers."):
                continue
            nested_dict[k.removeprefix(".subparsers.")] = v
            arg_dict.pop(k)
        arg_dict[subparser_field_name] = nested_type(**nested_dict)
    return d(**arg_dict)


@dataclass
class SubCommand1:
    glug: int = option(default=0)

    subparser_name: ClassVar[str] = "sub1"


@dataclass
class SubCommand2:
    garg: int = option(default=2)

    subparser_name: ClassVar[str] = "sub2"

    ex1: ClassVar[Exclusive] = exclusive_group()

    foo2: int | None = option(exclusive_group=ex1)
    bar2: int | None = option(exclusive_group=ex1)

@dataclass
class MyArgs:
    first_arg: int = positional()
    second_arg: list[str] | None = option()

    ex1: ClassVar[Exclusive] = exclusive_group()

    foo: int | None = option(exclusive_group=ex1, help="this is forwarded back to argparse")
    bar: int | None = option(exclusive_group=ex1, metavar="so_is_this")

    sub: SubCommand1 | SubCommand2 | None = subparsers(default=None)


print(parse_args(MyArgs))


# Most common use cases:
# - named arguments - typically optional, but can force required by not giving default
# - positional arguments
# - generally, if default is provided assume it is not required - otherwise it is, but allow override
# - type, help
# - store true, store false
# - store const
# - choices - via Literal/Enum
# - nargs - suppport *, + and N via some kind of list type annotation. ? support unclear for now.
# - short args: pass as an argument to annotation?
# -
# - support mutually exclusive groups - define a struct that inherits from our base class. All fields must
#   be Optional, and exactly one gets populated
# - support subparsers - again, define a nested struct