Python迭代器的实现原理

from typing import Iterable
print(
    isinstance("", Iterable),
    isinstance((), Iterable),
    isinstance([], Iterable),
    isinstance({}, Iterable),
    isinstance(set(), Iterable),
)  # True True True True True

print(
    isinstance(0, Iterable),
    isinstance(0.0, Iterable),
)  # False False

class A:
    def __getitem__(self, item):
        return f"参数item: {item}"
a = A()
#内部定义了 __getitem__
#首先可以让实例对象像字典一样访问属性
print(a["name"])  # 参数item: name
print(a["satori"])  # 参数item: satori

# 此外还可以像可迭代对象一样被for循环
# 循环的时候会自动给item传值，0 1 2 3...
# 如果内部出现了StopIteration，循环结束
# 否则会一直循环下去。这里我们手动break
for idx, val in enumerate(a):
    print(val)
    if idx == 5:
        break
"""
参数item: 0
参数item: 1
参数item: 2
参数item: 3
参数item: 4
参数item: 5
"""

from typing import Iterable
print(isinstance(a, Iterable))  # False

lst = [1, 2, 3]
#底层调用的其实是list.__iter__(lst)
#或者说PyList_Type.tp_iter(lst)
it = lst.__iter__()
print(it)  # <list_iterator object at 0x000001DC6E898640>
print(
    str.__iter__("")
)  # <str_iterator object at 0x000001DC911B8070>
print(
    tuple.__iter__(())
)  # <tuple_iterator object at 0x000001DC911B8070>

lst = [1, 2, 3]
# 等价于 type(lst).__iter__(lst)
it = iter(lst)

val = 0
def foo():
    global val
    val += 1
    return val
# iter可以接收一个参数: iter(可迭代对象)
# iter也可以接收两个参数: iter(可调用对象, value)
for i in iter(foo, 5):
    print(i)
"""
1
2
3
4
"""

static PyObject *
builtin_iter(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
{
    PyObject *v;
  
    // iter函数要么接收一个参数, 要么接收两个参数
    if (!_PyArg_CheckPositional("iter", nargs, 1, 2))
        return NULL;
    v = args[0];
    //如果接收一个参数
    //那么直接使用 PyObject_GetIter 获取对应的迭代器即可
    //可迭代对象的类型不同，那么得到的迭代器也不同
    if (nargs == 1)
        return PyObject_GetIter(v);
    // 如果接收的不是一个参数, 那么一定是两个参数
    // 如果是两个参数, 那么第一个参数一定是可调用对象
    if (!PyCallable_Check(v)) {
        PyErr_SetString(PyExc_TypeError,
                        "iter(v, w): v must be callable");
        return NULL;
    }
    // 获取value(哨兵)
    PyObject *sentinel = args[1];
    //调用PyCallIter_New
    //得到一个可调用的迭代器, calliterobject 对象
    
    return PyCallIter_New(v, sentinel);
}

PyObject *
PyObject_GetIter(PyObject *o)
{  
    //获取可迭代对象的类型对象
    PyTypeObject *t = Py_TYPE(o);
    //我们说类型对象定义的操作，决定了实例对象的行为
    //实例对象调用的那些方法都是定义在类型对象里面的
    //还是那句话：obj.func()等价于type(obj).func(obj)
    getiterfunc f;
    //所以这里是获取类型对象的tp_iter成员
    //也就是Python中的 __iter__
    f = t->tp_iter;
    //如果 f 为 NULL
    //说明该类型对象内部的tp_iter成员被初始化为NULL
    //即内部没有定义 __iter__ 
    //像str、tuple、list等类型对象，它们的tp_iter成员都是不为NULL的
    if (f == NULL) {
      //如果 tp_iter 为 NULL，那么解释器会退而求其次
      //检测该类型对象中是否定义了 __getitem__
      //如果定义了，那么直接调用PySeqIter_New
      //得到一个seqiterobject对象
      //下面的PySequence_Check负责检测类型对象是否实现了__getitem__
      //__getitem__ 对应 tp_as_sequence->sq_item
        if (PySequence_Check(o))
            return PySeqIter_New(o);
        // 走到这里说明该类型对象既没有__iter__、也没有__getitem__
        // 因此它的实例对象不具备可迭代的性质，于是抛出异常
        return type_error("'%.200s' object is not iterable", o);
    }
    else {
        // 否则说明定义了__iter__，于是直接进行调用
        // Py_TYPE(o)->tp_iter(o) 返回对应的迭代器
        PyObject *res = (*f)(o);
        // 但如果返回值res不为NULL、并且还不是迭代器
        // 证明 __iter__ 的返回值有问题，于是抛出异常
        if (res != NULL && !PyIter_Check(res)) {
            PyErr_FORMat(PyExc_TypeError,
                         "iter() returned non-iterator "
                         "of type '%.100s'",
                         Py_TYPE(res)->tp_name);
            Py_DECREF(res);
            res = NULL;
        }
        // 返回 res
        return res;
    }
}

typedef struct {
    PyObject_HEAD
    Py_ssize_t it_index;
    //指向创建该迭代器的列表
    PyListObject *it_seq;
} listiterobject;

from ctypes import *

class PyObject(Structure):
    _fields_ = [
        ("ob_refcnt", c_ssize_t),
        ("ob_size", c_void_p)
    ]

class ListIterObject(PyObject):
    _fields_ = [
        ("it_index", c_ssize_t),
        ("it_seq", POINTER(PyObject))
    ]

it = iter([1, 2, 3])
it_obj = ListIterObject.from_address(id(it))

# 初始的时候，索引为0
print(it_obj.it_index)  # 0
# 进行迭代
next(it)
# 索引自增1，此时it_index等于1
print(it_obj.it_index)  # 1
# 再次迭代
next(it)
# 此时it_index等于2
print(it_obj.it_index)  # 2
# 再次迭代
next(it)
# 此时it_index等于3
print(it_obj.it_index)  # 3

static PyObject *
builtin_next(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
{
    PyObject *it, *res;
  
    // 同样接收一个参数或者两个参数
    // 因为调用next函数时，可以传入一个默认值
    // 表示当迭代器没有元素可以迭代的时候，会返回指定的默认值
    if (!_PyArg_CheckPositional("next", nargs, 1, 2))
        return NULL;

    it = args[0];
    //第一个参数必须是一个迭代器
    if (!PyIter_Check(it)) {
        //否则的话, 抛出TypeError
        //表示第一个参数传递的不是一个迭代器
        PyErr_Format(PyExc_TypeError,
            "'%.200s' object is not an iterator",
            it->ob_type->tp_name);
        return NULL;
    } 
    //it->ob_type表示获取类型对象，也就是该迭代器的类型
    //可能是列表的迭代器、元组的迭代器、字符串的迭代器等等
    //具体是哪一种不重要，因为实现了多态
    //然后再获取tp_iternext成员，相当于__next__
    //拿到函数指针之后，传入迭代器进行调用
    res = (*it->ob_type->tp_iternext)(it);
    
    // 如果 res 不为 NULL, 那么证明迭代到值了, 直接返回
    if (res != NULL) {
        return res;
    } else if (nargs > 1) {
        //否则的话，说明 res == NULL，也就是有可能出错了
        //那么看nargs是否大于1, 如果大于1, 说明设置了默认值
        PyObject *def = args[1];
        // 如果出现异常
        if (PyErr_Occurred()) {
        // 那么就看该异常是不是迭代完毕时所产生的StopIteration异常
            if(!PyErr_ExceptionMatches(PyExc_StopIteration))
            // 如果不是，说明Python程序的逻辑有问题
            // 于是直接return NULL，结束执行
            // 然后在 Python 里面我们会看到打印到stderr中的异常信息
                return NULL;
            // 如果是 StopIteration，证明迭代完毕了
            // 但我们设置了默认值，那么就应该返回默认值
            // 而不应该抛出 StopIteration，于是将异常回溯栈给清空
            PyErr_Clear();
        }
        // 然后增加默认值的引用计数, 并返回
        Py_INCREF(def);
        return def;
    } else if (PyErr_Occurred()) {
        //走到这里说明 res == NULL，并且没有指定默认值
        //那么当发生异常时，将异常直接抛出
        return NULL;
    } else {
        // 都不是的话，直接抛出 StopIteration
        PyErr_SetNone(PyExc_StopIteration);
        return NULL;
    }
}

lst = [1, 2, 3]
it = iter(lst)
# 然后迭代，等价于next(it)
print(type(it).__next__(it))  # 1
print(type(it).__next__(it))  # 2
print(type(it).__next__(it))  # 3
# 但是next可以指定默认值
# 如果不指定默认值，或者还是type(it).__next__(it)
# 那么就会报错，会抛出StopIteration
print(next(it, 666))  # 666

//迭代器的类型对象
PyTypeObject PyListIter_Type = {
    PyVarObject_HEAD_INIT(&PyType_Type, 0)
    "list_iterator",                            
    sizeof(listiterobject),                     
    0,                                          
    
    (destructor)listiter_dealloc,               
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    0,                                          
    PyObject_GenericGetAttr,                    
    0,                                          
    0,                                          
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
    0,                                          
    (traverseproc)listiter_traverse,            
    0,                                          
    0,                                          
    0,                                          
    PyObject_SelfIter,                          
    (iternextfunc)listiter_next,                
    listiter_methods,                           
    0,                                          
};

static PyObject *
listiter_next(listiterobject *it)
{
    PyListObject *seq;  //列表
    PyObject *item;     //元素
    assert(it != NULL);
    //拿到具体对应的列表
    seq = it->it_seq;
    //如果seq为NULL，证明迭代器已经迭代完毕
    //否则它不会为NULL
    if (seq == NULL)
        return NULL;
    assert(PyList_Check(seq));
    //如果索引小于列表的长度，证明尚未迭代完毕
    if (it->it_index < PyList_GET_SIZE(seq)) {
      //通过索引获取指定元素
        item = PyList_GET_ITEM(seq, it->it_index);
      //it_index自增1
        ++it->it_index;
      //增加引用计数后返回
        Py_INCREF(item);
        return item;
    }
    //否则的话，说明此次索引正好已经超出最大范围
    //意味着迭代完毕了，将it_seq设置为NULL
    //并减少它的引用计数，然后返回
    it->it_seq = NULL;
    Py_DECREF(seq);
    return NULL;
}