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:

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!

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.