def get_snake_map()

in source/BattlesnakeGym/battlesnake_gym/snake.py [0:0]

• 35 lines of code
• 13 McCabe index (conditional complexity)

    def get_snake_map(self, return_type="Binary"):
        '''
        Return an image including the positions of the snakes

        Parameter:
        ----------
        return_type: string
            if Binary, a binary image is returned
            if Colour, an image based on the snake's colour is returned
            if Numbered, an image with 1 as the head, 2, 3, 4 as the body
                is returned

        Returns:
        --------
        map_image, np.array(self.map_size)
            image of the position of this snake
        '''        
        if return_type == "Colour":
            map_image = np.zeros((self.map_size[0], self.map_size[1], 3))
        else:
            map_image = np.zeros((self.map_size[0], self.map_size[1]))

        if not self._is_alive or self.is_head_outside_map():
            # To check if the snake is dead or not
            return map_image

        for i, location in enumerate(self.locations):
            if return_type == "Colour":
                map_image[location[0], location[1], :] = self.colour
            elif return_type == "Binary":
                map_image[location[0], location[1]] = 1
            elif return_type == "Numbered":
                map_image[location[0], location[1]] = i+1

        # Color the head differently
        if return_type == "Colour":
            map_image[self.get_head()[0], self.get_head()[1], :] *= 0.5
        elif return_type == "Binary":
            map_image[self.get_head()[0], self.get_head()[1]] = 5

        return map_image