通過argparse的numpy.array索引-如何正確執行？

Question

我想實現一個小的模塊，該模塊從HDF5文件加載數據，選擇特定的數據集，並對每個子集進行一些計算。 我決定使用argparse選項來選擇子集。 有沒有比使用字符串和手動解析更好的方法了？

下面的代碼應該可以實現我想要的功能，但是它不是很漂亮，並且還不能處理“：-3”之類的輸入。 一定會有更好的辦法！

謝謝！

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

我也傾向於解析字符串，但是利用slice來獲取slice。

更新 （請參閱下面的評論）

遵循以下內容：

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'))

結果：

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)

section = slice(*section)可能是個技巧，但如果您仔細看，它應該不難理解（我這樣做主要是因為，如果為我節省了一個if elif else系列，它可以檢查section的長度。

這不適用於多維索引； 在這種情況下，請先使用逗號分割。 如果要允許省略號，則必須分別捕獲三個點，並返回Ellipsis對象而不是切片。