aoc2016d8.py

import itertools

SCREEN_WIDTH = 50
SCREEN_HEIGHT = 6

def mask(size):
    return (1 << size) - 1

def iter_bits(num):
    while True:
        yield num & 1
        num >>= 1

def swap_x_y(screen):
    for col in zip(*map(iter_bits, screen)):
        new_row = 0
        for shift, bit in enumerate(col):
            new_row |= bit << shift
        yield new_row

def rect(screen, width, height):
    screen = list(screen)
    for index, row in enumerate(screen[:height]):
        screen[index] = row | mask(width)
    return screen

def rotate_row(screen, row, by, screen_width=SCREEN_WIDTH):
    screen = list(screen)
    screen[row] = (screen[row] << by) % mask(screen_width)
    return screen

def rotate_col(screen, col, by, screen_width=SCREEN_WIDTH, screen_height=SCREEN_HEIGHT):
    return list(
        itertools.islice(swap_x_y(
            rotate_row(
                itertools.islice(swap_x_y(screen), screen_width),
                col, by, screen_height
            )
        ), screen_height)
    )

def print_screen(screen, screen_width=SCREEN_WIDTH):
    for row in screen:
        print("".join("#" if potato else "_" for potato in itertools.islice(iter_bits(row), screen_width)))

def c_rect(screen, spec):
    return rect(screen, *map(int, spec.split("x")))

rot_sub_commands = {
    "row": rotate_row,
    "column": rotate_col
}

def c_rotate(screen, which, where_spec, _, by):
    return rot_sub_commands[which](screen, int(where_spec.split("=")[-1]), int(by))

commands = {
    "rect": c_rect,
    "rotate": c_rotate
}

def run_command(screen, command):
    argv = command.split()
    return commands[argv[0]](screen, *argv[1:])

if __name__ == "__main__":
    import fileinput
    screen = [0] * SCREEN_HEIGHT

    for line in fileinput.input():
        screen = run_command(screen, line)

    print_screen(screen)
    print("Total lit:", sum(sum(itertools.islice(iter_bits(row), SCREEN_WIDTH)) for row in screen))