micropython-samples/parse2d/parse2d.py

73 wiersze
2.9 KiB
Python

# parse2d.py Parse args for item access dunder methods for a 2D array.
# Released under the MIT License (MIT). See LICENSE.
# Copyright (c) 2023 Peter Hinch
# Called from __getitem__ or __setitem__ args is a 1-tuple. The single item may be an int or a
# slice for 1D access. Or it may be a 2-tuple for 2D access. Items in the 2-tuple may be ints
# or slices in any combination.
# As a generator it returns offsets into the underlying 1D array or list.
def do_args(args, nrows, ncols):
# Given a slice and a maximum address return start and stop addresses (or None on error)
# Step value must be 1, hence does not support start > stop (used with step < 0)
def do_slice(sli, nbytes):
step = sli.step if sli.step is not None else 1
start = sli.start if sli.start is not None else 0
stop = sli.stop if sli.stop is not None else nbytes
start = min(start if start >= 0 else max(nbytes + start, 0), nbytes)
stop = min(stop if stop >= 0 else max(nbytes + stop, 0), nbytes)
ok = (start < stop and step > 0) or (start > stop and step < 0)
return (start, stop, step) if ok else None # Caller should check
def ivalid(n, nmax): # Validate an integer arg, handle -ve args
n = n if n >= 0 else nmax + n
if n < 0 or n > nmax - 1:
raise IndexError("Index out of range")
return n
def fail(n):
raise IndexError("Invalid index", n)
ncells = nrows * ncols
n = args[0]
if isinstance(n, int): # Index into 1D array
yield ivalid(n, ncells)
elif isinstance(n, slice): # Slice of 1D array
cells = do_slice(n, ncells)
if cells is not None:
for cell in range(*cells):
yield cell
elif isinstance(n, tuple): # or isinstance(n, list) old versions of grid
if len(n) != 2:
fail(n)
row = n[0] # May be slice
if isinstance(row, int):
row = ivalid(row, nrows)
col = n[1]
if isinstance(col, int):
col = ivalid(col, ncols)
if isinstance(row, int) and isinstance(col, int):
yield row * ncols + col
elif isinstance(row, slice) and isinstance(col, int):
rows = do_slice(row, nrows)
if rows is not None:
for row in range(*rows):
yield row * ncols + col
elif isinstance(row, int) and isinstance(col, slice):
cols = do_slice(col, ncols)
if cols is not None:
for col in range(*cols):
yield row * ncols + col
elif isinstance(row, slice) and isinstance(col, slice):
rows = do_slice(row, nrows)
cols = do_slice(col, ncols)
if cols is not None and rows is not None:
for row in range(*rows):
for col in range(*cols):
yield row * ncols + col
else:
fail(n)
else:
fail(n)