Play with the A* algorithm in Pythonista

Astar.py

from scene import *
import heapq


TILE_W = 32
TILE_H = 32
COST_DIAG = 7
COST_ORTH = 5
MODE_DRAW = 1
MODE_START = 2
MODE_GOAL = 3
MODE_FIND = 4
MODE_FOUND = 5


class Astar(Scene):
    
    def setup(self):
        height = int(self.size.h - 56) / TILE_H
        width = int(self.size.w / TILE_W)
        self.map = list()
        for i in range(height):
            self.map.append([0]*width)
        self.green = (0.6, 0.8, 0.6)
        self.red = (0.8, 0.6, 0.6)
        self.blue = (0.6, 0.6, 0.8)
        self.colors = [(1,1,1),(0,0,0)]
        self.mode = None
        self.start = None
        self.goal = None
        self.path = None
    
    def get_xy(self, location):
        return int(location.x/TILE_W), int(location.y/TILE_H)
    
    def tile_get(self, x, y):
        try:
            return self.map[y][x]
        except IndexError:
            return None
    
    def tile_set(self, x, y, c):
        try:
            self.map[y][x] = c
        except IndexError:
            pass
    
    def draw_tile(self, x, y, color=None, border=1):
        if color:
            fill(*color)
        rect(x*TILE_W + border, y*TILE_H + border, TILE_W - border * 2, TILE_H - border * 2)
    
    def draw(self):
        background(0.9, 0.9, 0.9)
        for row, line in enumerate(self.map):
            for col, tile in enumerate(line):
                self.draw_tile(col, row, self.colors[tile])
        if self.mode == MODE_FIND or self.mode == MODE_FOUND:
            fill(0.6, 0.8, 0.8)
            for node in self.visited:
                self.draw_tile(node[0], node[1])
        tint(1, 1, 0.8)
        fill(*self.green)
        rect(8, self.size.h - 48, 80, 40)
        text('START', x=48, y=self.size.h - 28)
        if self.start:
            self.draw_tile(self.start[0], self.start[1], border=8)
        fill(*self.red)
        rect(104, self.size.h - 48, 80, 40)
        text('GOAL', x=144, y=self.size.h - 28)
        if self.goal:
            self.draw_tile(self.goal[0], self.goal[1], border=8)
        tint(1, 1, 1)
        fill(*self.blue)
        rect(200, self.size.h - 48, 80, 40)
        text('RUN', x=240, y=self.size.h - 28)
        if self.mode == MODE_FIND:
            self.find()
        if self.mode == MODE_FOUND:
            fill(0, 0.4, 0)
            for x, y in self.path:
                self.draw_tile(x, y, border=10)
    
    def find(self):
        estimate, score, node, parent = heapq.heappop(self.openqueue)
        if self.came_from.get(node) is not None:
            return
        self.came_from[node] = parent
        if node == self.goal:
            self.mode = MODE_FOUND
            self.reconstruct()
            return
        for cost, neighbor in self.neighbors(node):
            if cost < self.queued_cost.get(neighbor, 9999):
                self.queued_cost[neighbor] = cost
                self.visited.add(neighbor)
                heapq.heappush(self.openqueue, (score + cost + self.estimate(neighbor), score + cost, neighbor, node))
    
    def neighbors(self, node):
        for dy in range(-1, 2):
            for dx in range(-1, 2):
                if dx or dy:
                    x = node[0] + dx
                    y = node[1] + dy
                    if dx and dy:
                        cost = COST_DIAG
                    else:
                        cost = COST_ORTH
                    if self.tile_get(x, y) is 0:
                        yield (cost, (x, y))
    
    def estimate(self, node):
        dx = abs(node[0] - self.goal[0])
        dy = abs(node[1] - self.goal[1])
        ddiag = min(dx, dy)
        dorth = max(dx, dy) - ddiag
        return ddiag * COST_DIAG + dorth * COST_ORTH
    
    def reconstruct(self):
        node = self.goal
        self.path = list()
        while node != self.start:
            self.path.append(node)
            node = self.came_from[node]
    
    def touch_began(self, touch):
        x, y = self.get_xy(touch.location)
        col = self.tile_get(x, y)
        if col is None:
            if x < 3:
                self.mode = MODE_START
            elif x < 6:
                self.mode = MODE_GOAL
            elif x < 9 and self.start and self.goal:
                self.mode = MODE_FIND
                self.visited = set([self.start])
                self.came_from = dict()
                self.queued_cost = dict()
                self.openqueue = list()
                heapq.heappush(self.openqueue, (self.estimate(self.start), 0, self.start, None))
        else:
            if self.mode == MODE_START:
                self.start = (x, y)
                self.mode = None
            elif self.mode == MODE_GOAL:
                self.goal = (x, y)
                self.mode = None
            elif self.mode == None:
                self.mode = MODE_DRAW
                self.draw_col = 1-col
                self.tile_set(x, y, self.draw_col)
            elif self.mode == MODE_FOUND:
                self.mode = None
    
    def touch_moved(self, touch):
        if self.mode == MODE_DRAW:
            x, y = self.get_xy(touch.location)
            self.tile_set(x, y, self.draw_col)

    def touch_ended(self, touch):
        if self.mode == MODE_DRAW:
            self.mode = None

run(Astar(), frame_interval=4)