snakeq/snake.py

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


class Snake:
    """
    Programmable Game of Snake written in PyGame
    """

    def __init__(self, margin=80,length=4,grid_width=30,grid_height=30, grid_pts=30,fps=12):
        # Init attributes
        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
        # Setup pygame
        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):
        """
        Reset game state
        """
        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),scale=1):
        """
        Draw element on the snake area
        """
        delta=int(self.grid_pts-int(self.grid_pts*scale))
        rect=pygame.Rect(self.grid_pts*x+delta, self.grid_pts*y+self.margin+delta, self.grid_pts-2*delta, self.grid_pts-2*delta)
        pygame.draw.rect(self.screen,color,rect, 0)

    def draw_infos(self,color=(255,255,255),thickness=10):
        """
        Draw the information bar
        """
        # Draw separator
        rect=pygame.Rect(0, self.margin-thickness, self.grid_width*self.grid_pts, thickness)
        pygame.draw.rect(self.screen,(180,180,180),rect, 0)
        # Draw score
        text = self.font.render('Score: '+str(self.score)+"       Length: "+str(len(self.snake))+'       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, text_center)
    
    def new_apple(self):
        """
        Create a new apple
        """
        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):
        """
        Move the snake
        """
        # 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):
        """
        Draw the snake with color effects
        """
        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)
            if i>0:
                self.draw_pts(elt[0],elt[1],color=(color[0],int(color[1]/2),color[2]),scale=max(0.7,1-i/4))
            else:
                self.draw_pts(elt[0],elt[1],color=(color[0],int(color[1]/2),color[2]),scale=0.86)

    def has_loose(self):
        """
        Return true if the game is loose
        """
        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):
        """
        Play a game (one attempt)
        """
        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=(210,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
            # Check if an event handler is available
            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))