Search code examples
pythonpython-3.xpython-class

Class with auto-generated unique incremental integer values for attributes (the same as 'enum.auto()' do)


I need Enum-like style for a class to auto-generate vales. I can not use Enum because I need mutability, in other words, I will need to add attributes at runtime.

Pseudo-code:

def auto(attribute_name: str) -> int:
    # Calculate a value for an attribute...
    # E.g. the first attribute will have "1", the second will have "2", etc.
    return value


class Something:
    foo = auto()  # Must be auto-assigned as 1
    bar = auto()  # Must be auto-assigned as 2
    spam = auto()  # Must be auto-assigned as 3
    eggs = auto()  # Must be auto-assigned as 4


assert Something.foo == 1
assert Something.bar == 2
assert Something.spam == 3
assert Something.eggs == 4

Something.add("new_attribute")
assert Something.new_attribute == 5

I want to make sure that I will not reinvent the wheel by writing a lot of custom code. Are the any common ways to solve the issue?


Solution

  • The enum.auto is not actually a counter. It merely produces a sentinel value that is converted to a counter or similar by the enum.Enum machinery.

    While there is no directly re-useable, public variant of this mechanism, it is simple enough to write your own:

    auto = object()
    
    
    class AttributeCounter:
        """Baseclass for classes enumerating `auto` attributes"""
        # called when creating a subclass via `class Name(AttributeCounter):`
        def __init_subclass__(cls, **kwargs):
            cls._counter = 0
            # search through the class and count up `auto` attributes
            for name, value in cls.__dict__.items():
                if value is auto:
                    cls.add(name)
    
        @classmethod
        def add(cls, name: str):
            """Add attribute `name` with the next counter value"""
            count = cls._counter = cls._counter + 1
            setattr(cls, name, count)
    

    The auto value is an arbitrary sentinel for "counted attributes"; if desired one could turn it into a function or even re-use enum.auto.

    Using __init_subclass__ is a lightweight alternative to the metaclass magic used by enum. Similar to enum's use of _generate_next_value_, it provides the general machinery but allows overwriting the logic by replacing a single method add.


    Similar to enum.Enum, the counter functionality is added to your own types by inheritance:

    class Something(AttributeCounter):
        foo = auto
        bar = auto
        spam = auto
        eggs = auto
    
    
    assert Something.foo == 1
    assert Something.bar == 2
    assert Something.spam == 3
    assert Something.eggs == 4
    
    Something.add("new_attribute")
    assert Something.new_attribute == 5