Python Pygame - Snake 360

import pygame
from pygame import Vector2 as V
from pygame import Rect as R
from pygame.locals import *
from itertools import product
from random import randint
import time

SPATIAL_GRID_SIZE = 32

class spatial_hash:
    def __init__(s):
        s.grid = {}

    def insert(s, rect, entity):
        cells = []
        for cell in s._rect_cells(rect):
            cells.append(cell)
            items = s.grid.get(cell)
            if items is None:
                s.grid[cell] = [entity]
            else:
                items.append(entity)
        return cells

    def _rect_cells(s, rect):
        x1, y1 = rect.topleft
        x1 //= SPATIAL_GRID_SIZE
        y1 //= SPATIAL_GRID_SIZE
        x2, y2 = rect.bottomright
        x2 = x2 // SPATIAL_GRID_SIZE + 1
        y2 = y2 // SPATIAL_GRID_SIZE + 1
        return product(range(x1, x2), range(y1, y2))

    def query(s, rect):
        items = []
        for cell in s._rect_cells(rect):
            items.extend(s.grid.get(cell, ()))
        return items

class segment:
    def __init__(s, pos, size, id):
        s.pos = pos
        s.rect = R(s.pos - (size, size), (size * 2, size * 2))
        s.grid_refs = SH.insert(s.rect, s)
        s.id = id

class worm:
    INIT_DIRECTION = V(0, -1)
    INIT_LENGTH = 200
    TURNING_SPEED = 3
    GROW_TIME = 1

    def __init__(s):
        s.dead = False

        s.pos = V(PDR.center) - (0, s.INIT_LENGTH // 2)
        s.dv = V(s.INIT_DIRECTION)

        s.id = s.INIT_LENGTH
        s.safe_collision_distance = int(V().distance_to((SEGMENT_SIZE, SEGMENT_SIZE))) + SEGMENT_SIZE + 3

        s.body = [segment(s.pos // 1, SEGMENT_SIZE, s.id)]

        s.grow_timestamp = None

        for y_off in range(1, s.INIT_LENGTH):
            pos = s.body[0].pos + (0, y_off)
            s.body.append(segment(pos, SEGMENT_SIZE, s.id + y_off))
            pygame.draw.circle(PDS, WHITE, pos, SEGMENT_SIZE)

    def update(s):
        global FOOD

        s.pos += s.dv * TIME_DELTA

        if s.pos.x < -SEGMENT_SIZE: s.pos.x = PDR.w + SEGMENT_SIZE
        if s.pos.x > PDR.w + SEGMENT_SIZE: s.pos.x = -SEGMENT_SIZE
        if s.pos.y < -SEGMENT_SIZE: s.pos.y = PDR.h + SEGMENT_SIZE
        if s.pos.y > PDR.h + SEGMENT_SIZE: s.pos.y = -SEGMENT_SIZE

        ipos = s.pos // 1
        if ipos != s.body[0].pos:
            s.id += 1
            s.body.insert(0, segment(ipos, SEGMENT_SIZE, s.id))

            now = time.perf_counter()
            if s.grow_timestamp == None:
                tip = s.body.pop()
                for gr in tip.grid_refs:
                    SH.grid[gr].remove(tip)
                pygame.draw.circle(PDS, BLACK, tip.pos, SEGMENT_SIZE)
                pygame.draw.circle(PDS, WHITE, s.body[-1].pos, SEGMENT_SIZE)
            else:
                if now - (now if not s.grow_timestamp else s.grow_timestamp) >= s.GROW_TIME:
                    s.grow_timestamp = None

            pygame.draw.circle(PDS, WHITE, ipos, SEGMENT_SIZE)

            potention_collision_segments = SH.query(s.body[0].rect)
            for seg in potention_collision_segments:
                if seg.id <= s.body[0].id - s.safe_collision_distance:
                    if s.body[0].pos.distance_to(seg.pos) <= SEGMENT_SIZE * 2:
                        s.dead = True

            if s.body[0].pos.distance_to(FOOD.pos) <= FOOD_SIZE + SEGMENT_SIZE:
                FOOD.remove_food()
                FOOD = food()
                s.grow_timestamp = time.perf_counter()

        k = pygame.key.get_pressed()
        if k[K_LEFT]:
            s.dv.rotate_ip(-s.TURNING_SPEED * TIME_DELTA)
        elif k[K_RIGHT]:
            s.dv.rotate_ip(s.TURNING_SPEED * TIME_DELTA)

class food:
    def __init__(s):
        while True:
            s.pos = V(randint(FOOD_SIZE, PDR.w - FOOD_SIZE), randint(FOOD_SIZE, PDR.h - FOOD_SIZE))
            rect = R(s.pos - (FOOD_SIZE, FOOD_SIZE), (FOOD_SIZE * 2, FOOD_SIZE * 2))
            space_found = True
            for seg in SH.query(rect):
                if seg.pos.distance_to(s.pos) <= FOOD_SIZE + SEGMENT_SIZE:
                    space_found = False
                    break
            if space_found:
                break
        pygame.draw.circle(PDS, BLUE, s.pos, FOOD_SIZE)

    def remove_food(s):
        pygame.draw.circle(PDS, BLACK, s.pos, FOOD_SIZE)



BLACK = (64, 64, 64)
WHITE = (255, 255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

PDR = R(0, 0, 640, 480)
FPS = 120

pygame.init()
PDS = pygame.display.set_mode(PDR.size, FULLSCREEN | SCALED)
PDS.fill(BLACK)

SEGMENT_SIZE = 5
FOOD_SIZE = 10

SH = spatial_hash()
WORM = worm()
FOOD = food()

delta_timestamp = None
exit_demo = False

while not exit_demo:
    now = time.perf_counter()
    TIME_DELTA = (now - (delta_timestamp if delta_timestamp != None else now)) * FPS
    delta_timestamp = now

    events = pygame.event.get()
    for event in events:
        if event.type == KEYUP:
            if event.key == K_ESCAPE:
                exit_demo = True

    WORM.update()
    if WORM.dead == True:
        exit_demo = True

    pygame.display.update()

pygame.quit()