Class that tracks data of all its active instantiations?

I have a class Foo with its instances having a "balance" attribute. I'm designing it in such a way that Foo can track all the balances of its active instances. By active I mean instances that are currently assigned to a declared variable, of part of a List that is a declared variable.

a = Foo(50) # Track this
b = [ Foo(20) for _ in range(5) ] # Track this
Foo(20)     # Not assigned to any variable. Do not track this.

Another feature of Foo is that is has an overloaded "add" operator, where you can add two Foo's balances together or add to a Foo's balance by adding it with an int or float.

Example:

x = Foo(200)
x = x + 50
y = x + Foo(30)

Here is my code so far:

from typing import List

class Foo:

    foo_active_instances: List = []

    def __init__(self, balance: float = 0):
        Foo.foo_active_instances.append(self)
        self.local_balance: float = balance

    @property
    def balance(self):
        """
        The balance of only this instance.
        """
        return self.local_balance

    def __add__(self, addend):
        """
        Overloading the add operator
        so we can add Foo instances together.
        We can also add more to a Foo's balance
        by just passing a float/int
        """
        if isinstance(addend, Foo):
            return Foo(self.local_balance + addend.local_balance)
        elif isinstance(addend, float | int):
            return Foo(self.local_balance + addend)
        
    @classmethod
    @property
    def global_balance(cls):
        """
        Sum up balance of all active Foo instances.
        """
        return sum([instance.balance for instance in Foo.foo_active_instances])

But my code has several issues. One problem is when I try to add a balance to an already existing instance, like:

x = Foo(200)
x = x + 50  # Problem: This instantiates another Foo with 200 balance.
y = Foo(100)

# Expected result is 350, because 250 + 100 = 350.
    
# Result is 550
# even though we just added 50 to x.
print(Foo.global_balance)

Another problem is replacing a Foo instance with None doesn't remove it from Foo.foo_active_instances.

k = Foo(125)
k = None
    
# Expected global balance is 0, 
# but the balance of the now non-existing Foo still persists
# So result is 125.
print(Foo.global_balance) 

I tried to make an internal method that loops through foo_active_instances and counts how many references an instance has. The method then pops the instance from foo_active_instance if it doesn't have enough. This is very inefficient because it's a loop and it's called each time a Foo instance is made and when the add operator is used.

How do I rethink my approach? Is there a design pattern just for this problem? I'm all out of ideas.


The weakref module is perfect for this design pattern. Instead of making foo_active_instances a list, you can make it a weakref.WeakSet. This way, when a Foo object's reference count falls to zero (e.g., because it wasn't bound to a variable), it will be automatically removed from the set.

class Foo:
   foo_active_instances = weakref.WeakSet()

   def __init__(self, balance: float = 0) -> None:
       Foo.foo_active_instances.add(self)
       ...

In order to add Foo objects to a set, you'll have to make them hashable. Maybe something like

class Foo:
    ...

    def __hash__(self) -> int:
        return hash(self.local_balance)

You can use inspect to check if the __init__ or __add__ methods have been called as part of an assignment statement. Additionally, you can keep a default parameter in __init__ to prevent increasing your global sum by the value passed to it when creating a new Foo object from __add__:

import inspect, re
def from_assignment(frame):
   return re.findall('[^\=]\=[^\=]', inspect.getframeinfo(frame).code_context[0])

class Foo:
   global_balance = 0
   def __init__(self, balance, block=False):
      if not block and from_assignment(inspect.currentframe().f_back):
         Foo.global_balance += balance
      self.local_balance = balance
   def __add__(self, obj):
      if from_assignment(inspect.currentframe().f_back) and not hasattr(obj, 'local_balance'):
         Foo.global_balance += obj
      return Foo(getattr(obj, 'local_balance', obj), True)

a = Foo(50)
b = [Foo(20) for _ in range(5)]
Foo(20)
print(Foo.global_balance) #150

x = Foo(200)
x = x + 50
y = Foo(100)
print(Foo.global_balance) #350