jeremiahbuddha
jeremiahbuddha

Reputation: 10191

Specifying types and patterns using argparse choices

I have several different input files, saved with case numbers ".1", ".2", etc to differentiate them. I am building a script that operates on these files, and I would like to use argparse to allow the user to specify the particular case number to operate on, or use "_" to specify the last saved case (i.e. the input file with the largest case number). Something like;

> ls 
   file.1, file.2, file.3
> my_script.py 2
   (operates on file.2)
> my_script.py _
   (operate on file.3)

Is there a way I can specify "any integer" as one choice and "_" as the second choice? Something like;

parser = argparse.ArgumentParser()
parser.add_argument('case', choices=[anyint, '_'])

Upvotes: 1

Views: 3658

Answers (1)

srgerg
srgerg

Reputation: 19339

You could use the type argument to add_argument(...) instead. For example:

import os
import argparse

def intOrUnderscore(s):
    if s != '_':
        return int(s)
    cases = (n for n in os.listdir(".") if n.startswith("file."))
    return max(int(c[c.rindex(".")+1:]) for c in cases)

parser = argparse.ArgumentParser()
parser.add_argument('case', type=intOrUnderscore)

args = parser.parse_args()
print args.case

When I run this I get:

$ ls
file.1  file.2  file.3  s.py
$ python s.py 2
2
$ python s.py _
3

Alternately, you could build the choices list in code:

import os
import argparse

cases = [n[n.rindex(".")+1:] for n in os.listdir(".") if n.startswith("file.")]
cases.append("_")
parser = argparse.ArgumentParser()
parser.add_argument('case', choices = cases)

args = parser.parse_args()
print args.case

Upvotes: 7

Related Questions