Python to generate pieces of Soma cube puzzle (WIP)

soma.py

DIMENSIONS = 2
SCALE = 3
OUTER_SCALE = SCALE + 2
STEPS = tuple(OUTER_SCALE**x for x in reversed(range(DIMENSIONS)))

def generate_all_cells():
    def do_slice(offset, depth):
        if depth == DIMENSIONS - 1:
            for i in range(1, SCALE + 1):
                yield offset + i
        else:
            for i in range(1, SCALE + 1):
                subslice_offset = offset + i * STEPS[depth]
                yield from do_slice(subslice_offset, depth + 1)

    yield from do_slice(0, 0)

def is_valid_I_piece(p):
    c1, c2, c3 = p

    for step in STEPS:
        if c1 + step == c2 and c2 + step == c3:
            return True

    return False

def is_valid_L_piece(p):
    c1, c2, c3 = p

    num_found = 0
    for step in STEPS:
        if c1 + step == c2 or c1 + step == c3 or c2 + step == c3:
            num_found += 1
            if num_found == 2:
                return True
            continue

    return False

# 20 21 22 23 24
# 15 16 17 18 19
# 10 11 12 13 14
#  5  6  7  8  9
#  0  1  2  3  4

def generate_all_pieces(cells):
    print(cells)
    for i,c1 in enumerate(cells):
        for j,c2 in enumerate(cells[i+1:], start=i+1):
            for k,c3 in enumerate(cells[j+1:], start=j+1):
                p = (c1, c2, c3)
                if is_valid_I_piece(p) or is_valid_L_piece(p):
                    yield p

ALL_CELLS = list(generate_all_cells())

ALL_PIECES = list(generate_all_pieces(ALL_CELLS))

print(len(ALL_PIECES))
for p in ALL_PIECES:
    print(p, is_valid_I_piece(p), is_valid_L_piece(p))



is_valid_L_piece((6, 7, 12))