numpy.array indices via argparse - How to do it properly?

Question

I want to implement a small module that loads data from a HDF5 file, selects a specific data set and does some calculation on each of the subsets. I decided to use an argparse option for selecting the subsets. Is there a better way than using a string and manually parsing it?

The code below should do what I want but it isn't very pretty and is not able to handle inputs like ':-3' yet. There must be a better way!

Thanks!

parser.add_argument('--filename', required=True, 
               help='the file to open')
parser.add_argument('--snapshots', default="-1", type=str, 
               help="select snapshots (default: %(default)s) ")
[...]
H5 = tables.openFile(filename)

snapshot = args.snapshots
maxval = len(H5.root.trajectory[:])
if snapshot ==":" or snapshot.lower() == "all":
    snapshot = np.arange(maxval)
elif ':' in snapshot:
    if snapshot.index(':') == 0:
        snapshot = np.arange(int(snapshot[1:]))
    elif snapshot.index(':') == len(snapshot) - 1:
        snapshot = np.arange(int(snapshot[:-1]),maxval)
    else:
        from_to = snapshot.split(':')
        snapshot = np.arange(int(from_to[0]), int(from_to[1]))
else:
    snapshot = int(snapshot)
selected = H5.root.trajectory[snapshot]

Answer 1

I tend to parse a string as well, but make use of slice to obtain the slice.

Updated (see comments below)

Something along the following lines::

def parsesection(snapshot):
    if snapshot == 'all':
        return slice(None)
    try:
        section = int(snapshot)
    except ValueError:
        section = [int(s) if s else None for s in snapshot.split(':')]
        if len(section) > 3:
            raise ValueError('snapshots input incorrect')
        section = slice(*section)
    return section

print(parsesection(':5'))
print(parsesection('5:15'))
print(parsesection('5:15:3'))
print(parsesection('3'))
print(parsesection('5::3'))
print(parsesection('::2'))
print(parsesection('all'))

Results:

slice(None, 5, None)
slice(5, 15, None)
slice(5, 15, 3)
3
slice(5, None, 3)
slice(None, None, 2)
slice(None, None, None)

The section = slice(*section) might feel as a trick, but if you look carefully, it shouldn't be hard to understand (I do it mostly because if saves me an if elif else series that checks for the length of section .

This will not work for multi-dimensional indexing; split on a comma first in that case. And in case you want to allow for ellipsis, you'll have to catch the three dots separately, and return an Ellipsis object instead of a slice.