KD Tree in Pygame

kd_tree.py

import pygame, random

class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

class Kd_tree:

    class Axis:
        X = 0
        Y = 1
        DIMS = 2

    class Leaf:

        def __init__(self, point):
            self.point = point
        
        def draw(self, surface, box):
            pygame.draw.circle(surface, [0,0,0], [self.point.x, self.point.y], 4)

    class Node:

        def __init__(self, axis, point_list, parent = None):
            self.axis = axis
            self.value = None
            self.parent = parent
            mid = None
            if self.axis is Kd_tree.Axis.X:
                point_list = sorted(point_list, key=lambda p: p.x)
            else:
                point_list = sorted(point_list, key=lambda p : p.y)

            mid = len(point_list) // 2

            if len(point_list) > 2:
                if self.axis is Kd_tree.Axis.X:
                    self.value = point_list[mid-1].x
                else:
                    self.value = point_list[mid-1].y
                self.left = Kd_tree.Node((axis+1) % Kd_tree.Axis.DIMS, point_list[:mid], self)
                self.right = Kd_tree.Node((axis+1) % Kd_tree.Axis.DIMS, point_list[mid:], self)
            else:
                self.left = Kd_tree.Leaf(point_list[0])
                self.right = None
                
                if self.axis is Kd_tree.Axis.X:
                    self.value = point_list[0].x
                else:
                    self.value = point_list[0].y

                if len(point_list) == 2:
                    self.right = Kd_tree.Leaf(point_list[1])
                    
        
        def draw(self, surface, box):
            if self.axis is Kd_tree.Axis.X:
                pygame.draw.line(surface, [255,0,0], [self.value, box[0][1]], [self.value, box[1][1]])
                self.left.draw(surface, [
                    box[0],
                    [self.value, box[1][1]]
                ])
                if self.right:
                    self.right.draw(surface, [
                        [self.value, box[0][1]],
                        box[1]
                    ])
            else:
                pygame.draw.line(surface, [0,255,0], [box[0][0], self.value], [box[1][0], self.value])
                self.left.draw(surface, [
                    box[0],
                    [box[0][0], self.value]
                ])
                if self.right:
                    self.right.draw(surface, [
                        [box[0][0], self.value],
                        box[1]
                    ])

    def __init__(self, point_list):
        self.root = self.Node(Kd_tree.Axis.X, point_list)

    def draw(self, surface, box):
        self.root.draw(surface, box)

points = [Point(random.randint(0,800), random.randint(0,600)) for i in range(10)]
tree = Kd_tree(points)

pygame.init()

surface = pygame.display.set_mode([800, 600])

while True:
    surface.fill([255,255,255])
    tree.draw(surface, [[0,0], [800,600]])
    pygame.display.flip()