pythonpython-3.xenumerate

Overriding enumerate for custom class


I have a custom class that's essentially a list, but with negative indices being valid indices, rather than referring to elements from the rear of the list.

from collections.abc import Sequence


class MultiFloorPlan(Sequence):
    def __init__(self):
        super().__init__()
        self._floors = []
        self._subfloors = []

    def __eq__(self, other):
        if not isinstance(other, MultiFloorPlan):
            return NotImplemented
        return self._subfloors == other._subfloors and self._floors == other._floors

    def _reindex(self, floor):
        if floor >= 0:
            return self._floors, floor
        return self._subfloors, -floor - 1

    def __len__(self):
        return len(self._subfloors) + len(self._floors)

    def __getitem__(self, floor):
        floor_list, floor = self._reindex(floor)
        return floor_list[floor]

    def __delitem__(self, floor):
        floor_list, floor = self._reindex(floor)
        del floor_list[floor]

    def __iter__(self):
        for plan in self._subfloors:
            yield plan
        for plan in self._floors:
            yield plan

    def __reversed__(self):
        for plan in reversed(self._floors):
            yield plan
        for plan in reversed(self._subfloors):
            yield plan

    def __contains__(self, value):
        return value in self._subfloors or value in self._floors

    def append(self, subfloor=False):
        if subfloor:
            return self._subfloors.append(None)  # For this example we append a dummy None
        return self._floors.append(None)         # value instead of an actual Plan instance

Is it possible to get the built-in enumerate to return the floor indices, rather than the values shifted to the non-negative integers? Example:

mfp = MultiFloorPlan()
for _ in range(5):
    mfp.append(subfloor=False)
    mfp.append(subfloor=True)

for floor, _ in enumerate(mfp):
    print(floor)
# This prints 0 1 2 3 4 5 6 7 8 9, but I'd like it to print -5 -4 -3 -2 -1 0 1 2 3 4

Solution

  • The enumerate builtin function operates on arbitrary iterators, so it cannot behave in a special way for your class. However, a clever user of enumerate can produce the results you want, they just need to set it up properly by passing in an appropriate starting index. You might provide that setup in the class, with a method:

    def enumerate(self):
        return enumerate(self, -len(self._subfloors))