Fix to handle flags in argparse