SZI2019SmieciarzWmi/DataModels/Container.py

class Container():
    def __init__(self, max, y=0, g=0, b=0):
        self.yellow, self.green, self.blue = y, g, b
        self.max = max

    def empty(self, trash_type=None, trash=None):
        if trash_type == None:
            self.yellow, self.green, self.blue = 0, 0, 0
        elif trash_type == "yellow":
            trash[0] += self.yellow
            self.yellow = 0
        elif trash_type == "green":
            trash[1] += self.green
            self.green = 0
        elif trash_type == "blue":
            trash[2] += self.blue
            self.blue = 0
        return trash


    def status(self):
        return [self.yellow, self.green, self.blue]


    def add(self, trash):
        my_trash = [self.yellow, self.green, self.blue]
        leftovers = [0, 0, 0]
        for i in range(0, len(trash)):
            while(my_trash[i] < self.max and trash[i] > 0):
                my_trash[i] += 1
                trash[i] -= 1
        self.yellow, self.green, self.blue = my_trash
        result = [0, 0, 0]
        for i in range(0, len(trash)):
            result[i] = trash[i] - leftovers[i]

        return result