如何为自定义MultiFloorPlan类重写enumerate以返回自定义索引？

我实现了一个继承自Sequence的自定义类MultiFloorPlan，它将非负索引对应地上层（_floors），负索引对应专属地下层（_subfloors），而非常规列表的末尾元素。但使用内置enumerate遍历实例时，返回的是从0开始的连续非负索引，我希望它能返回对应实际楼层的索引（例如：地下5层到地上4层对应索引-5到4）。

原类实现代码：

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)  # 示例中用None代替实际Plan实例
        return self._floors.append(None)

解决方案 #

内置enumerate会生成连续的整数索引（默认从0开始），无法感知自定义类的特殊索引逻辑，因此需要通过以下两种方式实现需求：

方式一：自定义enumerate生成器方法 #

在MultiFloorPlan类中添加自定义的enumerate方法，直接生成对应楼层的索引和元素：

def enumerate(self):
    # 生成地下层索引：从 -len(_subfloors) 到 -1
    for idx, plan in enumerate(self._subfloors):
        yield -(len(self._subfloors) - idx), plan
    # 生成地上层索引：从0到 len(_floors)-1
    for idx, plan in enumerate(self._floors):
        yield idx, plan

使用示例：

# 创建实例并添加5个地下层、5个地上层
plan = MultiFloorPlan()
for _ in range(5):
    plan.append(subfloor=True)
for _ in range(5):
    plan.append()

# 使用自定义enumerate遍历
for floor_idx, item in plan.enumerate():
    print(floor_idx, item)
# 输出：-5 None, -4 None, ..., -1 None, 0 None, 1 None, ..., 4 None

方式二：手动结合zip与索引序列 #

如果不想修改类，可以在遍历前手动构造正确的索引序列，通过zip结合元素遍历：

plan = MultiFloorPlan()
# 同上添加楼层...

# 构造完整索引序列
subfloor_indices = list(range(-len(plan._subfloors), 0))
floor_indices = list(range(len(plan._floors)))
full_indices = subfloor_indices + floor_indices

# 遍历输出
for idx, item in zip(full_indices, plan):
    print(idx, item)

内容的提问来源于stack exchange，提问作者Mate de Vita