Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Get grouped structure using argparse.ArgumentParser()

I'd like to translate something like:

test.py --group 1 --opt1 foo1 --opt2 bar1 \
        --group 2 --opt1 foo2 \
        --group 3 --opt1 foo3 --opt2 bar3 --opt3 baz3

into something like a dictionary (like Namespace):

{
  "1": {"opt1": "foo1", 
        "opt2": "bar1"},
  "2": {"opt1": "foo2"}, 
  "3": {"opt1": "foo3", 
        "opt2": "bar3",
        "opt3": "baz3"}
}
  • The input can be another format if needed.
  • I want to raise an error if other than (opt1, opt2, opt3) is used
  • I should use argparse.ArgumentParser().

Can you help?

like image 440
Tetsuro Nakamura Avatar asked Sep 10 '26 09:09

Tetsuro Nakamura


1 Answers

Seemed like a cool problem -- argparse lets you extend its callbacks in two ways, either with type=... for custom types or action=... with custom behaviour. I don't think the problem as stated is possible with type= so I went with a solution involving action=

The basic approach to writing a custom action involves inheriting from argparse.Action and overriding __call__ (and optionally __init__)

Here's a start at implementing your idea, you'll likely need to extend / change it (for instance, it might make sense to call super().__call__(...) in the action classes to have the base behaviour handled (for types, etc.) -- I just went with the simplest thing that worked to satisfy your question.

import argparse
import pprint


class GroupAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        group, = values
        namespace._current_group = group
        groups = namespace.__dict__.setdefault('groups', {})
        if group in groups:
            raise argparse.ArgumentError(self, f'already specified: {group}')
        groups[group] = {}


class AppendToGroup(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        value, = values
        if not getattr(namespace, '_current_group', None):
            raise argparse.ArgumentError(self, 'outside of group!')
        namespace.groups[namespace._current_group][self.dest] = value


parser = argparse.ArgumentParser()
parser.add_argument('--group', action=GroupAction)
parser.add_argument('--opt1', action=AppendToGroup)
parser.add_argument('--opt2', action=AppendToGroup)
parser.add_argument('--opt3', action=AppendToGroup)
pprint.pprint(parser.parse_args().groups)

Note here that I use the namespace to store state, so unfortunately at the end you'll end up with a little extra args._current_group.

The __dict__.setdefault is a bit sneaky so I'll explain it -- it's a slightly shorter way to write:

if not hasattr(namespace, 'groups'):
    namespace.groups = {}
groups = namespace.groups

The basic strategy I used was to store the current group, and append to that group when seeing other arguments

Here's the sample in action:

$ python3.7 t.py --group 1 --opt1 a --opt2 b --opt3 c
{'1': {'opt1': 'a', 'opt2': 'b', 'opt3': 'c'}}
$ python3.7 t.py --group 1 --opt1 a --opt2 b --opt3 c --group 2
{'1': {'opt1': 'a', 'opt2': 'b', 'opt3': 'c'}, '2': {}}
$ python3.7 t.py --group 1 --opt1 a --opt2 b --opt3 c --group 2 --opt3 c
{'1': {'opt1': 'a', 'opt2': 'b', 'opt3': 'c'}, '2': {'opt3': 'c'}}
$ python3.7 t.py --group 1 --group 1
usage: t.py [-h] [--group GROUP] [--opt1 OPT1] [--opt2 OPT2] [--opt3 OPT3]
t.py: error: argument --group: already specified: 1
$ python3.7 t.py --opt1 a --group 1
usage: t.py [-h] [--group GROUP] [--opt1 OPT1] [--opt2 OPT2] [--opt3 OPT3]
t.py: error: argument --opt1: outside of group!
like image 158
Anthony Sottile Avatar answered Sep 13 '26 01:09

Anthony Sottile