CustomClassCollection迭代时元素类型提示失效，求解决方案

问题重现 #

使用CustomClassCollection迭代时，IDE无法识别元素的CustomClass类型，导致类型提示失效：

foo = CustomClass("abc")
bar = CustomClass("cba")

collection = CustomClassCollection([foo, bar])

for item in collection:
    item.<无类型提示>

当前CustomClassCollection的迭代相关实现存在明显错误：

def __next__(self) -> CustomClass:
    yield from self._collection.values()

def __iter__(self) -> CustomClass:
    return next(self)

问题根源 #

1. __iter__返回类型错误：迭代器协议要求__iter__返回迭代器对象，而非单个元素，当前标注的返回类型CustomClass完全不符合要求。
2. 迭代逻辑混乱：
   • __next__是迭代器的方法，负责单次返回元素并终止，不能用yield from（这是生成器函数的写法）。
   • __iter__里调用next(self)会直接返回第一个元素，而非返回迭代器本身，彻底违反迭代器协议。

修复方案 #

方案一：复用现有迭代器（最简最优） #

直接返回self._collection.values()的迭代器，同时修正类型提示：

from collections.abc import Iterator

class CustomClassCollection(KeyIndexedCollection, JsonSerializable):
    # 保留其他原有方法，仅修改/删除以下内容
    def __iter__(self) -> Iterator[CustomClass]:
        return iter(self._collection.values())
    
    # 删除错误的__next__方法，不需要自己实现

方案二：自定义迭代器实现（仅需特殊逻辑时用） #

如果需要自定义迭代行为，正确实现迭代器协议：

from collections.abc import Iterator
from typing import Optional

class CustomClassCollection(KeyIndexedCollection, JsonSerializable):
    def __init__(self, items: Optional[List[CustomClass]] = None):
        super().__init__(items)
        self._collection = {}
        self._add_items_impl(self._collection, take_with_default(items, []))
        self._iter_keys = None  # 迭代状态变量

    # 保留其他原有方法
    def __iter__(self) -> Iterator[CustomClass]:
        self._iter_keys = iter(self._collection.keys())
        return self

    def __next__(self) -> CustomClass:
        try:
            key = next(self._iter_keys)
            return self._collection[key]
        except StopIteration:
            self._iter_keys = None  # 重置状态支持多次迭代
            raise

额外修正 #

get方法的返回类型标注错误，应改为Optional[CustomClass]：

def get(self, key: str, default: Optional[Any] = None) -> Optional[CustomClass]:
    return self._collection.get(key, default)

验证 #

修复后，迭代时item会被正确识别为CustomClass，IDE会正常提示name()等方法：

for item in collection:
    item.name()  # 此处显示类型提示

内容的提问来源于stack exchange，提问作者cxnt