Here's a quick example of making a copy of the screen surface.
It calls window.copy() to make an "undo point" on start, and whenever the user clicks on [KEEP]. If the user clicks on [UNDO], the copied surface is re-blitted to the screen to erase the existing one.

I'm not sure if keeping a copy of the screen is a great way to do this, it uses a lot of memory. A better approach would be to store the drawing primitives in a list, then only re-draw up to N-1 entry on undo.
import pygame
# Window size
WINDOW_WIDTH = 400
WINDOW_HEIGHT = 400
WINDOW_SURFACE = pygame.HWSURFACE|pygame.DOUBLEBUF|pygame.RESIZABLE
DARK_BLUE = ( 3, 5, 54 )
YELLOW = ( 250, 250, 0 )
RED = ( 200, 0, 0 )
### initialisation
pygame.init()
pygame.mixer.init()
window = pygame.display.set_mode( ( WINDOW_WIDTH, WINDOW_HEIGHT ), WINDOW_SURFACE )
pygame.display.set_caption("Undo Example")
font = pygame.font.SysFont('freesansbold.ttf', 30)
undo_text = font.render( "[UNDO]", True, YELLOW )
undo_rect = undo_text.get_rect()
undo_rect.x = 50
keep_text = font.render( "[KEEP]", True, YELLOW )
keep_rect = keep_text.get_rect()
keep_rect.x = 200
window.fill( DARK_BLUE )
### The undo buffer
undo_save = None
### Main Loop
clock = pygame.time.Clock()
pen_down = False
last_point = None
done = False
while not done:
# Handle user-input
for event in pygame.event.get():
if ( event.type == pygame.QUIT ):
done = True
elif ( event.type == pygame.MOUSEMOTION ):
mouse_pos = event.pos
if ( pen_down ):
if ( last_point != None ):
pygame.draw.line( window, RED, last_point, mouse_pos )
last_point = mouse_pos
elif ( event.type == pygame.MOUSEBUTTONDOWN ):
pen_down = True
if ( last_point == None ):
last_point = event.pos
elif ( event.type == pygame.MOUSEBUTTONUP ):
pen_down = False
last_point = None
### Handle the KEEP/UNDO button presses
if ( undo_rect.collidepoint( event.pos ) ):
if ( undo_save != None ):
window.blit( undo_save, ( 0, 0 ) )
print( "UNDO Image Restored" )
if ( keep_rect.collidepoint( event.pos ) ):
undo_save = window.copy()
print( "UNDO Image Saved" )
window.blit( undo_text, undo_rect )
window.blit( keep_text, keep_rect )
# Update the window, but not more than 60fps
pygame.display.flip()
# make the first undo automatically
if ( undo_save == None ):
undo_save = window.copy()
# Clamp FPS
clock.tick(60)
pygame.quit()