How can I draw a square onto a grid when clicked with Pygame?

Question

I'm using the code from Python: How to make drawn elements snap to grid in pygame to draw a grid and show a square that snaps onto the grid. The thing is that the square only stays there until you stop clicking, which is not what I want.

So, what I did was changing the code from this:

def draw(self, surface):
    click = pygame.mouse.get_pressed()
    if click[0]:
        pygame.draw.rect(surface, (255, 255, 255), self.square)

To this:

def draw(self, surface):
    pygame.draw.rect(surface, (255, 255, 255), self.square) #it draws the square without the need to keep the mouse button down
    click = pygame.mouse.get_pressed()
    if click[0]:
        pygame.draw.rect(surface, (255, 255, 255), self.square)

And this makes the square follow the cursor and snap onto the grid. But when clicked, it doesn't draw a square into the grid.

So my overall question is, how can I draw a square onto the grid when mouse clicked using the code from the link provided?