import argparse
import unittest

class MyArgumentParser(argparse.ArgumentParser):
    # could subclass argparse._ActionsContainer instead
    def add_titled_mutually_exclusive_group(self, **kwargs):
        # works only if there is a title
        group = self.add_argument_group(**kwargs)
        group1 = group.add_mutually_exclusive_group() # no kwargs
        return group1

    def add_titled_mutually_exclusive_group(self, **kwargs):
        # works with and without a title
        if 'title' in kwargs:
            container = self.add_argument_group(**kwargs)
            # remove title for next step
            kwargs.pop('title', None)
            kwargs.pop('description', None)
        else:
            container = self
        group = argparse._MutuallyExclusiveGroup(container, **kwargs)
        container._mutually_exclusive_groups.append(group)
        return group

MUTEX_HELP_HAS_NO_TITLE = """usage: PROG_CMD [-h] [--arg1 ARG1 | --arg2 ARG2]

Program description

optional arguments:
  -h, --help   show this help message and exit
  --arg1 ARG1  arg1 help
  --arg2 ARG2  arg2 help

Epilog
"""


MUTEX_HELP_INCLUDES_TITLE = """usage: PROG_CMD [-h] [--arg1 ARG1 | --arg2 ARG2]

Program description

optional arguments:
  -h, --help   show this help message and exit

Group title:
  Group description

  --arg1 ARG1  arg1 help
  --arg2 ARG2  arg2 help

Epilog
"""


class EntitledMutexTestCase(unittest.TestCase):

    def test_mutex_without_title_or_description(self):

        parser = MyArgumentParser(
            prog='PROG_CMD',
            description='Program description',
            epilog='Epilog')

        mutex_group = parser.add_titled_mutually_exclusive_group()
        mutex_group.add_argument('--arg1', help='arg1 help')
        mutex_group.add_argument('--arg2', help='arg2 help')

        help_string = parser.format_help()
        self.assertEqual(help_string, MUTEX_HELP_HAS_NO_TITLE)

    def test_mutex_with_title_and_description(self):

        parser = MyArgumentParser(
            prog='PROG_CMD',
            description='Program description',
            epilog='Epilog')

        mutex_group = parser.add_titled_mutually_exclusive_group(
            title='Group title',
            description='Group description')
        mutex_group.add_argument('--arg1', help='arg1 help')
        mutex_group.add_argument('--arg2', help='arg2 help')

        help_string = parser.format_help()
        self.assertEqual(help_string, MUTEX_HELP_INCLUDES_TITLE)


if __name__ == "__main__":
    unittest.main()