snakeq/snake.py

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


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

    def __init__(self, startat=(0,0), margin=80,length=4,grid_width=30,grid_height=30, grid_pts=30,fps=15):
        # 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
        self.startat=startat
        self.last_score=-1
        self.score=0
        # 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))
        self.clock = pygame.time.Clock()
        # Start game
        self.new_game()
        self.draw()

    def new_game(self):
        """
        Reset game state
        """
        self.snake=[self.startat]*self.default_length
        self.direction=3 # Like clock (12=up, 3=right, 6=bottom, 9=left)
        self.new_apple()
        self.last_score=self.score
        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))
        if self.apple in self.snake:
            self.new_apple()

    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 draw_apple(self,radius=None):
        """ 
        Draw the apple on the screen
        """
        if radius==None:
            radius=self.grid_pts/2
        pygame.draw.circle(self.screen, (210,0,0), (self.apple[0]*self.grid_pts+int(self.grid_pts/2),self.apple[1]*self.grid_pts+self.margin+int(self.grid_pts/2)), radius)
        pygame.draw.circle(self.screen, (170,0,0), (self.apple[0]*self.grid_pts+int(self.grid_pts/2),self.apple[1]*self.grid_pts+self.margin+int(self.grid_pts/2)), radius/1.5)


    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 draw(self):
        self.screen.fill((0,0,0))
        self.draw_snake()
        self.draw_apple()
        self.draw_infos()
        pygame.display.flip()

    def play(self,direction,handle_quit=True):
        """
        Play using wall clock directions (12=up, 3=right, 6=down and 9=left)
        """
        # Play
        if  (self.direction==12 and direction!=6) or\
            (self.direction==6 and direction!=12) or\
            (self.direction==3 and direction!=9) or\
            (self.direction==9 and direction!=3):
            self.direction=direction
        self.move()
        # Return code
        code=0
        if self.apple==self.snake[0]:
            self.snake.append(self.snake[len(self.snake)-1])
            self.new_apple()
            self.score+=1
            code=1
        elif self.has_loose():
            self.new_game()
            code=-1
        # Refresh screen
        self.draw()
        self.clock.tick(self.fps)
        # Ensure we not quit
        if handle_quit:
            for event in pygame.event.get():
                    if event.type == pygame.QUIT: 
                        pygame.quit()
                        sys.exit()
        return(code)
    
    def play2(self,direction):
        """
        Play using ascii directions
        """
        if direction.lower()=="up":
            return(self.play(self,12))
        elif direction.lower()=="right":
            return(self.play(self,3))
        elif direction.lower()=="down":
            return(self.play(self,6))
        elif direction.lower()=="left":
            return(self.play(self,9))
    
    def play3(self,direction):
        """
        Play using 0123 directions (0=right, 1=down, 3=left and 4=up)
        """
        if direction == 0:
            return(self.play(3))
        elif direction == 1:
            return(self.play(6))
        elif direction == 2:
            return(self.play(9))
        elif direction == 3:
            return(self.play(12))

    def play_with_keyboard(self):
        """
        Play a game using keyboard
        """
        while True:
            # Check inputs
            for event in pygame.event.get():
                if event.type == pygame.QUIT: 
                    pygame.quit()
                    sys.exit()
                elif 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 self.play(self.direction,handle_quit=False) <0:
                break