能否实现根据参数返回游标或生成数据的方法？遇生成器属性错误

你的函数因为包含yield关键字，被Python认定为生成器函数。这类函数调用时不会立即执行函数体，而是直接返回一个生成器对象——哪怕你传入return_cursor=True，得到的也不是游标，而是生成器。只有当你迭代这个生成器时，函数体才会开始执行，但此时return cur的结果不会直接返回给调用者，而是被封装在StopIteration异常中，这就导致你访问rowcount时触发AttributeError。

你的需求完全可以实现，以下是几种推荐的调整方式：

方案1：拆分函数（推荐，职责单一） #

把游标获取和结果生成拆成两个独立函数，逻辑更清晰：

def get_query_cursor(connection):
    query = """SELECT ... """
    cur = connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
    cur.execute(query)
    return cur

def fetch_rows(connection):
    cur = get_query_cursor(connection)
    info(f"Found {cur.rowcount} record(s).")
    for row in cur:
        yield row

• 需要游标时，直接调用get_query_cursor(conn)
• 需要迭代结果时，调用fetch_rows(conn)

方案2：单函数返回列表（牺牲内存换简洁） #

如果不需要生成器的懒加载特性，可以把结果转为列表返回，避免函数成为生成器：

def fetch_rows(connection, return_cursor=False):
    query = """SELECT ... """
    cur = connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
    cur.execute(query)

    if return_cursor:
        return cur
    else:
        info(f"Found {cur.rowcount} record(s).")
        return list(cur)

注意：结果集很大时，这种方式会占用较多内存。

方案3：给生成器绑定游标（保留懒加载） #

如果想同时保留生成器的内存优势和获取游标的能力，可以给生成器对象绑定游标属性：

def fetch_rows(connection, return_cursor=False):
    query = """SELECT ... """
    cur = connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
    cur.execute(query)

    if return_cursor:
        return cur
    else:
        info(f"Found {cur.rowcount} record(s).")
        def result_generator():
            for row in cur:
                yield row
        gen = result_generator()
        gen.cursor = cur  # 将游标绑定到生成器上
        return gen

调用时，迭代生成器获取数据，需要游标属性时直接访问gen.cursor.rowcount即可。

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