一文帶你解密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)
# 從 C 的角度上看，就是 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 物件
    /*
    該物件位於 Objects/iterobject.c 中
    */
    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__
        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))
            // 如果不是，說明程式的邏輯有問題
            // 於是直接return NULL，結束執行
            // 然後在 Python 裡面我們會看到列印到stderr中的異常資訊
                return NULL;
            // 如果是 StopIteration，證明迭代完畢了
            // 但我們設定了預設值，那麼就應該返回預設值
            // 而不應該丟擲 StopIteration，於是將異常回溯棧給清空
            PyErr_Clear();
        }
        // 然後增加預設值的參照計數, 並返回
        Py_INCREF(def);
        return def;
    } else if (PyErr_Occurred()) {
        //走到這裡說明程式出異常了，並且沒有指定預設值
        //那麼這種情況，不管什麼異常都直接丟擲
        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

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;
}