Pygame Simple Infinite State Machine

# This is a simple example.
import pygame

class State:
  def __init__(self, control):
    self.control = control

  def on_draw(self, surface): pass
  def on_event(self, event): pass
  def on_update(self, delta): pass

class StateControl:
  def __init__(self, caption, size, flags=0, fps=60):
    pygame.display.set_caption(caption)
    self.display = pygame.display.set_mode(size, flags)
    self.rect = self.display.get_rect()
    self.clock = pygame.time.Clock()
    self.delta = 0
    self.fps = fps
    self.running = True

    self.current_state = None
    self.next_state = None

  def main_loop(self):
    while self.running:
      if self.next_state:
        self.current_state = self.next_state
        self.next_state = None

      for event in pygame.event.get():
        if event.type != pygame.QUIT:
          self.current_state.on_event(event)
        else:
          self.running = False

      self.current_state.on_update(self.delta)
      self.current_state.on_draw(self.display)
      pygame.display.flip()
      self.delta = self.clock.tick(self.fps) / 1000

  def run(self, state):
    if state:
      self.current_state = state
      self.main_loop()

class RedState(State):
  def __init__(self, control):
    super().__init__(control)
    self.background = 'firebrick'

  def on_draw(self, surface):
    surface.fill(self.background)

  def on_event(self, event):
    if event.type == pygame.KEYDOWN:
      self.control.next_state = BlueState(self.control)

class BlueState(State):
  def __init__(self, control):
    super().__init__(control)
    self.background = 'dodgerblue'

  def on_draw(self, surface):
    surface.fill(self.background)

  def on_event(self, event):
    if event.type == pygame.KEYDOWN:
      self.control.next_state = RedState(self.control)

def main():
  control = StateControl("Example Infinite State Machine", (800, 600))
  control.run(RedState(control))

if __name__ == "__main__":
  pygame.init()
  main()
  pygame.quit()