多个列表理解 return 空列表与上下文管理器

Question

我正在读取一个压缩的 csv 文件，并且想在不使用 pandas 的情况下只提取特定的列。我当前的代码只有 returns 一个用于第一个列表理解的列表，但不用于以下列表。如何在使用上下文管理器时提取多列？

输入文件：

col1,col2,col3
1,2,3
a,b,c

我的代码

import gzip
import csv
import codecs

with gzip.open(r"myfile.csv.gz", "r") as f:
    content = csv.reader(codecs.iterdecode(f, "utf-8"))

    col_2 = [row[1] for row in content] # Returns [2, "b"]
    col_3 = [row[2] for row in content] # Returns []

预期输出：

col_2: [2, "b"]
col_3: [3, "c"]

Answer 1

问题不是上下文管理器引起的，而是只能读取一次的生成器引起的。

您可以使用 itertools.tee 复制它：

import gzip
import csv
import codecs

with gzip.open(r"myfile.csv.gz", "r") as f:
    content = csv.reader(codecs.iterdecode(f, "utf-8"))
    from itertools import tee

    c1, c2 = tee(content) # from now on, do not use content anymore
    
    col_2 = [row[1] for row in c1]
    col_3 = [row[2] for row in c2]

输出：

>>> col_2
['col2', '2', 'b']

>>> col_3
['col3', '3', 'c']

使用经典循环

不过，更好的方法是使用经典循环。这避免了必须循环两次值：

with gzip.open(r"myfile.csv.gz", "r") as f:
    content = csv.reader(codecs.iterdecode(f, "utf-8"))

    col_2 = []
    col_3 = []
    for row in content:
        col_2.append(row[1])
        col_3.append(row[2])

多个列表理解 return 空列表与上下文管理器

Multiple list comprehensions return empty list with context manager

python

csv

contextmanager

使用经典循环