argparse_fill_dict_action.py

import argparse
from typing import Any, Dict, List, Optional

class FillDict(argparse.Action):
    """Argparse action that fills in a dictionary with nargs.

    General usecase:

    Subclass this and set custom_keys and defaults (if there are any).
    Set that subclass as the action of an argparse argument, with nargs="+".
    Now it will be possible to pass in multiple args and turn them into a dictionary.
    Extending the dictionary is simple: just add more keys (+defaults). Note that args
    must be in proper order, with no gaps. Otherwise, improper parsing will occur. Thus,
    keys should be sorted based on "optionality", since an optional arg towards the
    beginning can't be skipped.
    """

    #: A list of keys that can be filled by this action
    custom_keys: List[str]

    #: A list of defaults for some keys
    defaults: Optional[Dict[str, Any]] = None

    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        self._nargs = nargs
        super(FillDict, self).__init__(option_strings, dest, nargs=nargs, **kwargs)

    @classmethod
    def max_nargs(cls) -> int:
        """Get how many arguments can go into this parser."""
        return len(cls.custom_keys)

    def _fill_defaults(self, parsed):
        """Fill in defaults that are missing."""
        if self.defaults is not None:
            for key, default in self.defaults.items():
                if key not in parsed:
                    parsed[key] = default
        return parsed

    def __call__(self, parser, namespace, values, option_string=None):
        """Parse the args.

        Args must be in proper order, with no gaps. Otherwise, improper parsing will occur.
        """
        assert (
            len(values) <= self.max_nargs()
        ), f"More values provided than requested. Values={values}"
        parsed = dict()
        for i, value in enumerate(values):
            parsed[self.custom_keys[i]] = value
        parsed = self._fill_defaults(parsed)
        setattr(namespace, self.dest, parsed)