snakeq/snake/snake.py

#!/usr/bin/env python
import sys, pygame, random


class Snake:

    def __init__(self, margin=80,length=10,grid_width=40,grid_height=40, grid_pts=20,fps=10):
        self.grid_width=grid_width
        self.grid_height=grid_height
        self.grid_pts=grid_pts
        self.margin=margin
        self.default_length=length
        self.attempt=0
        self.fps=fps
        pygame.init()
        self.font=pygame.font.SysFont(pygame.font.get_default_font(), int(self.margin/2))
        self.font_small=pygame.font.SysFont(pygame.font.get_default_font(), int(self.margin/2.5))
        self.screen=pygame.display.set_mode((grid_width*grid_pts,grid_height*grid_pts+margin))

    def new_game(self):
        self.snake=[(0,0)]*self.default_length
        self.direction=3 # Like clock (12=up, 3=right, 6=bottom, 9=left)
        self.new_apple()
        self.score=0
        self.attempt+=1

    def draw_pts(self,x,y,color=(255,255,255)):
        rect=pygame.Rect(self.grid_pts*x, self.grid_pts*y+self.margin, self.grid_pts, self.grid_pts)
        pygame.draw.rect(self.screen,color,rect, 0)

    def draw_infos(self,color=(255,255,255),thickness=5):
        rect=pygame.Rect(0, self.margin-thickness, self.grid_width*self.grid_pts, thickness)
        pygame.draw.rect(self.screen,color,rect, 0)
        text = self.font.render('Score '+str(self.score)+"     Length "+str(len(self.snake)), True, color)
        text_center=text.get_rect(center = (self.grid_width*self.grid_pts // 2, (self.margin-thickness) // 2))
        self.screen.blit(text, text_center)
        text = self.font_small.render('Attempt '+str(self.attempt), True, color)
        text_center=text.get_rect(center = (self.grid_width*self.grid_pts // 2, (self.margin-thickness) // 2))
        self.screen.blit(text, (self.grid_pts/2,text_center[1]))

    def new_apple(self):
        self.apple=(random.randint(0,self.grid_width-1),random.randint(0,self.grid_height-1))
        while self.apple in self.snake:
            self.apple=(random.randint(0,self.grid_width),random.randint(0,self.grid_height))

    def move(self):
        # Update tail
        if len(self.snake)>1:
            tmp=self.snake[0]
            for i in range(1,len(self.snake)):
                newtmp=self.snake[i]
                self.snake[i]=tmp
                tmp=newtmp
        # Update head
        h=self.snake[0] # Head
        if self.direction==3:
            self.snake[0]=(h[0]+1,h[1])
        elif self.direction==9:
            self.snake[0]=(h[0]-1,h[1])
        elif self.direction==12:
            self.snake[0]=(h[0],h[1]-1)
        else:
            self.snake[0]=(h[0],h[1]+1)

    def draw_snake(self):
        for i in range(0,len(self.snake)):
            color=(0,150,150) if i==0 else (0,max(255-i*10,120),0) 
            elt=self.snake[i]
            self.draw_pts(elt[0],elt[1],color=color)

    def has_loose(self):
        if self.snake.count(self.snake[0])>1:
            return(True)
        h=self.snake[0]
        if h[0]<0 or h[1]<0 or h[0] >= self.grid_width or h[1] >= self.grid_height:
            return(True)
        return(False)


    def run(self, event_handler=None):
        clock = pygame.time.Clock()
        ignore_has_loose=True
        self.new_game()
        while True:
            self.screen.fill((0,0,0))
            self.draw_snake()
            self.draw_pts(self.apple[0],self.apple[1],color=(255,0,0))
            self.draw_infos()
            # Check for loose
            if not(ignore_has_loose) and self.has_loose():
                break
            else:
                ignore_has_loose=False
            # Check inputs
            for event in pygame.event.get():
                if event.type == pygame.QUIT: 
                    pygame.quit()
                    sys.exit()
                elif event_handler==None and event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT and self.direction != 3:
                        self.direction=9
                        break
                    elif event.key == pygame.K_RIGHT and self.direction != 9:
                        self.direction=3
                        break
                    elif event.key == pygame.K_UP and self.direction != 6:
                        self.direction=12
                        break
                    elif event.key == pygame.K_DOWN and self.direction != 12:
                        self.direction=6
                        break
            if event_handler!=None:
                event_handler(self)
            self.move()
            # Check for eating apple
            if self.apple==self.snake[0]:
                self.snake.append(self.snake[len(self.snake)-1])
                self.new_apple()
                self.score+=1
            pygame.display.flip()
            clock.tick(self.fps)
        return(self.score)


game=Snake()

def event_handler(game):
    if game.snake[0][0]==10:
        game.direction=6

for i in range(0,10):
    score=game.run()
    print("Game ended with "+str(score))