Creating Buttons in PyGame

For starters GameState is missing TITLE value.

class GameState(Enum):
    # ...
    TITLE = 3

Adding this makes the code run.

The log_in() function does not handle the window being closed. You must handle the pygame.QUIT event, in every event loop. For example:

def log_in( screen ):
    # ...

    while True:
        mouse_up = False
        for event in pygame.event.get():
            if ( event.type == pygame.QUIT ):
                pygame.event.post( pygame.event.Event( pygame.QUIT ) )    # re-send the quit event to the next loop
                return GameState.QUIT
            elif ( event.type == pygame.MOUSEBUTTONUP and event.button == 1 ):
                mouse_up = True    # Mouse button 1 weas released

        ui_action = return_btn.update( pygame.mouse.get_pos(), mouse_up )
        if ui_action is not None:
            print( "log_in() - returning action" )
            return ui_action

        screen.fill(PINK)
        return_btn.draw(screen)
        pygame.display.flip()

The UIElement.update() looks like it should return the self.action when the mouse button is released over the control. However, in the existing code, nothing is ever returned. Probably it needs to be something like this:

class UIElement( Sprite ):
    # ...
    def update(self, mouse_pos, mouse_up):
        """ Track the mouse, setting the self.mouse_over.  Also check 
            if the mouse-button was clicked while over this control
            returning the pre-defined self.action, if so.     """

        result = None                    # No click => no action
        if self.rect.collidepoint(mouse_pos):
            self.mouse_over = True
            if ( mouse_up ):
                result = self.action     # Mouse was clicked on element, add action
        else:
            self.mouse_over = False
        return result

After these changes, your script runs OK before dropping into an outer loop when the button is clicked. The outer loop also does not handle exiting properly, but probably it's just the same set of changes over again.

It would be better to only have one user-input processing loop. having these separate event-loops is causing the same problem in multiple locations. Work out a way to have a single event processing function, then adapt your UI code to use it. This will make your code easier to write and debug going forward.

Ref: All Code

import pygame

import pygame.freetype

from pygame.sprite import Sprite

from pygame.rect import Rect

from enum import Enum

PINK = (250, 100, 100)

WHITE = (255, 255, 255)

BLACK = (0,0,0)

def create_surface_with_text(text, font_size, text_rgb, bg_rgb):

    font = pygame.freetype.SysFont("Arial", font_size, bold=True)

    surface, _ = font.render(text=text, fgcolor=text_rgb, bgcolor=bg_rgb)

    return surface.convert_alpha()


class UIElement(Sprite):
   
    def __init__(self, center_position, text, font_size, bg_rgb, text_rgb, action=None):
     
        self.mouse_over = False  

        # what happens when the mouse is not over the element
        default_image = create_surface_with_text(
            text=text, font_size=font_size, text_rgb=text_rgb, bg_rgb=bg_rgb
        )

        # what happens when the mouse is over the element
        highlighted_image = create_surface_with_text(
            text=text, font_size=font_size * 1.1, text_rgb=text_rgb, bg_rgb=bg_rgb
        )

        self.images = [default_image, highlighted_image]

        self.rects = [
            default_image.get_rect(center=center_position),
            highlighted_image.get_rect(center=center_position),
        ]

        self.action = action
        
        super().__init__()

    @property
    def image(self):
        return self.images[1] if self.mouse_over else self.images[0]

    @property
    def rect(self):
        return self.rects[1] if self.mouse_over else self.rects[0]

    def update(self, mouse_pos, mouse_up):
        """ Track the mouse, setting the self.mouse_over.  Also check 
            if the mouse-button was clicked while over this control
            returning the pre-defined self.action, if so.     """

        result = None                    # No click => no action
        if self.rect.collidepoint(mouse_pos):
            self.mouse_over = True
            if ( mouse_up ):
                result = self.action     # Mouse was clicked on element, add action
        else:
            self.mouse_over = False
        return result


    def draw(self, surface):
        surface.blit(self.image, self.rect)


def main():
    pygame.init()

    screen = pygame.display.set_mode((800, 600))
    game_state = GameState.LOGIN

    while True:
        if game_state == GameState.LOGIN:
            game_state = log_in(screen)

        if game_state == GameState.SIGNUP:
            game_state = sign_up(screen)

        if game_state == GameState.RETURN:
            game_state = title_screen(screen)

        if game_state == GameState.QUIT:
            pygame.quit()
            return

def title_screen(screen):

    login_btn = UIElement(
        center_position=(400,300),
        font_size=30,
        bg_rgb=WHITE,
        text_rgb=BLACK,
        text="Log In",
        action=GameState.LOGIN,
    )

    signup_btn = UIElement(
        center_position=(400,200),
        font_size=30,
        bg_rgb=WHITE,
        text_rgb=BLACK,
        text="Log In",
        action=GameState.LOGIN,
    )

    uielement = UIElement(
        center_position=(400, 100),
        font_size=40,
        bg_rgb=PINK,
        text_rgb=BLACK,
        text="Welcome to the Automated Timetable Program",
        action=GameState.QUIT,
     )

    buttons = [login_btn, signup_btn]

    while True:
        mouse_up = False
        for event in pygame.event.get():
            if event.type == pygame.MOUSEBUTTONUP and event.button == 1:
                mouse_up = True
            elif event.type == pygame.QUIT:
                pygame.quit()
                sys.exitIO
        screen.fill(PINK)

        for button in buttons:
            ui_action = button.update(pygame.mouse.get_pos(),mouse_up)
            if ui_action is not None:
                return ui_action
            button.draw(screen)

        pygame.display.flip()

def log_in(screen):

    return_btn = UIElement(
        center_position=(140, 570),
        font_size=20,
        bg_rgb=WHITE,
        text_rgb=BLACK,
        text="Return to main menu",
        action=GameState.TITLE,
    )

    while True:
        mouse_up = False
        for event in pygame.event.get():
            if ( event.type == pygame.QUIT ):
                pygame.event.post( pygame.event.Event( pygame.QUIT ) )    # re-send the quit event to the next loop
                return GameState.QUIT
            elif ( event.type == pygame.MOUSEBUTTONUP and event.button == 1 ):
                mouse_up = True    # Mouse button 1 weas released

        ui_action = return_btn.update( pygame.mouse.get_pos(), mouse_up )
        if ui_action is not None:
            print( "log_in() - returning action" )
            return ui_action

        screen.fill(PINK)
        return_btn.draw(screen)
        pygame.display.flip()



class GameState(Enum):

    LOGIN = -1
    SIGNUP = 0
    RETURN = 1
    QUIT = 2
    TITLE=3


if __name__ == "__main__":

    main()