Water jug puzzle simulation using Pygame

Question

I am making a puzzling tool using Pygame. We have unlimited containers that we can choose their maximum capacity, and how much water is in them. We can also choose the size of each unit for convenience.

Here is how it goes:

enter image description here

Here is my code:

import pygame
pygame.init()
wn = pygame.display.set_mode((600, 600))
class TextBox():
 def __init__(self, x, y, w, h, color_active=(0, 0, 200), color_inactive=(0, 0, 0), default='', pad=10, font=pygame.font.SysFont('Arial', 22)):
 self.pad = pad
 self.input_box = pygame.Rect(x+self.pad, y+self.pad, w, h)
 self.w = w
 self.h = h
 self.font = font
 self.color_inactive = color_inactive
 self.color_active = color_active
 self.color = self.color_inactive
 self.default = default
 self.text = default
 self.active = False
 def draw(self):
 if not self.text:
 self.text = '0'
 txt = self.font.render(self.text, True, self.color)
 width = max(self.w, txt.get_width()+self.pad)
 self.input_box.w = width
 wn.blit(txt, (self.input_box.x+5, self.input_box.y+3))
 pygame.draw.rect(wn, self.color, self.input_box, 2)
 def check_status(self, pos):
 if self.input_box.collidepoint(pos):
 self.active = not self.active
 else:
 self.active = False
 self.color = self.color_active if self.active else self.color_inactive
objs = []
class Obj():
 def __init__(self, color, x, y, w, h, default):
 self.w = w
 self.h = h
 self.color = color
 self.rect = pygame.rect.Rect(x, y, w, h)
 self.dragging = False
 self.pad = 10
 self.txt_box = TextBox(x, y, w-20, 30, pad=10, default='0')
 self.txt_box2 = TextBox(x, y-default, w-20, 30, pad=10, default='0')
 self.new = True
 self.straw = pygame.rect.Rect(x, y, self.pad, self.pad)
 self.line = False
 self.straw_start = None
 self.straw_end = None
 self.snap = False
 self.snap_y = 450
 objs.append(self)
 def clicked(self, pos):
 return self.rect.collidepoint(pos)
 def clicked_straw(self, pos):
 return self.straw.collidepoint(pos)
 def trans(self, o):
 empty = int(o.txt_box.text) - int(o.txt_box2.text)
 full = int(self.txt_box2.text)
 total = full if empty >= full else empty
 self.txt_box2.text = str(int(self.txt_box2.text) - total)
 o.txt_box2.text = str(int(o.txt_box2.text) + total)
 def offset_click(self, pos):
 self.dragging = True
 self.offset_x = self.rect.x - pos[0]
 self.offset_y = self.rect.y - pos[1]
 self.offset_x2 = self.txt_box.input_box.x - pos[0]
 self.offset_y2 = self.txt_box.input_box.y - pos[1]
 self.offset_x3 = self.txt_box2.input_box.x - pos[0]
 self.offset_y3 = self.txt_box2.input_box.y - pos[1]
 def offset_drag(self, pos):
 self.new = False
 self.rect.x = self.straw.x = self.offset_x + pos[0]
 self.rect.y = self.straw.y = self.offset_y + pos[1]
 self.txt_box.input_box.x = self.offset_x2 + pos[0]
 self.txt_box.input_box.y = self.offset_y2 + pos[1]
 self.txt_box2.input_box.x = self.offset_x3 + pos[0]
 self.txt_box2.input_box.y = self.offset_y3 + pos[1]
 
 def draw(self):
 pygame.draw.rect(wn, self.color, self.rect)
 amt1, amt2 = int(self.txt_box.text), int(self.txt_box2.text)
 bt_y = self.txt_box.input_box.bottom
 if amt1:
 w, h = self.rect.w, self.h * amt2
 x, y = self.rect.x, bt_y-h + self.pad - self.h
 pygame.draw.rect(wn, (145, 255, 255), (x, y, w, h))
 self.txt_box2.draw()
 else:
 self.txt_box2.text = '0'
 self.rect.h = self.h * (amt1 + 1)
 self.rect.y = bt_y - self.rect.h + self.pad
 self.txt_box.draw()
 pygame.draw.rect(wn, (0, 55, 255), self.straw)
 def draw_straw(self):
 if self.line:
 pygame.draw.line(wn, (0, 255, 255), self.straw_start, self.straw_end, 5)
num = TextBox(400, 20, 100, 50, color_inactive=(255, 255, 255), color_active=(200, 200, 255), default='50', font=pygame.font.SysFont('Arial', 40))
obj = Obj((255, 255, 255), 30, 30, 70, 50, int(num.default))
while True:
 for event in pygame.event.get():
 if event.type == pygame.QUIT:
 pygame.quit()
 elif event.type == pygame.MOUSEBUTTONDOWN:
 if event.button == 1:
 for o in objs:
 o.txt_box.check_status(event.pos)
 o.txt_box2.check_status(event.pos)
 if o.clicked_straw(event.pos):
 o.line = True
 o.straw_start = event.pos
 o.straw_end = event.pos
 elif o.clicked(event.pos):
 o.offset_click(event.pos)
 num.check_status(event.pos)
 elif event.button == 3:
 for o in objs:
 if o.clicked(event.pos) and not o.new:
 o.snap = not o.snap
 if o.snap:
 o.offset_click(event.pos)
 o.offset_drag((event.pos[0], o.snap_y+event.pos[1]-o.rect.y-int(o.txt_box.text)*o.h-int(num.text)))
 o.dragging = False
 
 elif event.type == pygame.MOUSEBUTTONUP:
 if event.button == 1:
 for o in objs:
 if o.line:
 o2 = [o for o in objs if o.clicked_straw(event.pos)]
 if o2:
 o.trans(o2[0])
 o.line = False
 o.dragging = False
 elif event.type == pygame.MOUSEMOTION:
 for o in objs:
 if o.line:
 o.straw_end = event.pos
 if o.dragging:
 if o.new:
 o.new = False
 obj = Obj((255, 255, 255), 30, 30, 70, int(num.text), int(num.default))
 o.offset_drag(event.pos)
 elif event.type == pygame.KEYDOWN:
 if num.active:
 if event.key == pygame.K_BACKSPACE:
 num.text = '0'
 elif event.unicode.isdigit():
 if num.text == '0':
 num.text = ''
 num.text += event.unicode
 for o in objs:
 o.h = int(num.text)
 o.rect.h = int(num.text)
 o.rect.y = o.txt_box.input_box.bottom-o.rect.h + o.pad
 o.straw.y = o.rect.y - o.h * int(o.txt_box.text)
 for o in objs:
 if o.txt_box.active:
 if event.key == pygame.K_BACKSPACE:
 o.txt_box.text = '0'
 o.rect.h = o.h
 elif event.unicode.isdigit():
 if o.txt_box.text == '0':
 o.txt_box.text = ''
 o.txt_box.text += event.unicode
 o.rect.h = o.h * (int(o.txt_box.text) + 1)
 o.rect.y = o.txt_box.input_box.bottom - o.rect.h + o.pad
 o.straw.y = o.rect.y
 elif o.txt_box2.active:
 if event.key == pygame.K_BACKSPACE:
 o.txt_box2.text = '0'
 elif event.unicode.isdigit():
 if int(o.txt_box2.text + event.unicode) <= int(o.txt_box.text):
 if o.txt_box2.text == '0':
 o.txt_box2.text = '' 
 o.txt_box2.text += event.unicode
 wn.fill((0, 100, 0))
 num.draw()
 for o in objs:
 o.draw()
 for o in objs:
 o.draw_straw()
 pygame.display.flip()

As you can see, my code is in one heck of a mess. Believe this is the opposite of DRY.

Can someone show me how to merge the redundant lines of code?

Also, there might be a bug or too swimming in my code, I don't know. Big thanks to those who find one!

• hjpotter92: how do you remove a placed bucket/jug?
• Chocolate: @hjpotter92 They are permanent buckets. Choose wisely.
• fartgeek: Cool program +1

Answer 1

This seems like a really interesting program. There are definitely some steps you could take to make the program easier to maintain and expand. Here are my thoughts.

• I would avoid causing side effects in the constructor of Obj. Having the line objs.append(self) means that whenever you want to create an Obj, you always need to have the global list objs initialised appropriately.

• You could abstract some common features out of Obj and TextBox as there is definitely some shared logic. Perhaps something like this:

class Clickable:
 def __init__(self, x, y, w, h, pad=0):
 self.x = x
 self.y = y
 self.w = w
 self.h = h
 self.pad = pad
 self.rect = pygame.Rect(x + pad, y + pad, w, h)
 
 def clicked(self, pos):
 return self.rect.collidepoint(pos)

Then, you can simplify TextBox, for example, by writing

class TextBox(Clickable):
 def __init__(self, x, y, w, h, color_active=(0, 0, 200), color_inactive=(0, 0, 0), default='', pad=10, font=pygame.font.SysFont('Arial', 22)):
 super(x, y, w, h, pad)
 # Rest of code here
 def check_status(self, pos):
 if self.clicked(pos):
 ...

• Your game loop seems quite tightly coupled to the implementation of the Obj class, and it manipulates quite a lot of the internal state of each Obj. Instead, perhaps it would make more sense to pass on each relevant event to an object, for example writing o.mouse_moved(...) and then letting the object itself decide how to react. As it stands, the main loop decides whether to change each object, which breaks the idea of "encapsulation".

As a rough guide of the direction you should try to go with the program:

• Decouple the main game loop and the internals of each Obj. Try to make it so the main loop knows as little as possible about the Obj itself, and instead, if a click or mouse move occurs, tell the object and let it decide what action is appropriate. This will make it easier if you want to add more classes like Obj, because you can implement general methods such as draw() and clicked() which the game loop can call.

• Share logic through inheritance. Where two objects behave in a similar way or you feel that you're repeating yourself, try and figure out what is the same about the two classes, then abstract that out into a parent class as I demonstrated above.

• Chocolate: Thanks! I'm still studying your post, but I've noticed the merging of the 2 for loops. The 2 loops are intentional so that the straws will never be covered by any of the jugs.
• htl: @user229550 Indeed, that is a mistake on my part, thanks. I've amended this now.