fileinput.hook_compressed 有时给我字符串，有时给我字节

Question

我正在尝试从多个文件中读取行。有些是 gzip 压缩的，有些是纯文本文件。在Python 2.7中，我一直在使用以下代码并且它有效：for line in fileinput.input(filenames, openhook=fileinput.hook_compressed):    match = REGEX.match(line)    if (match):        # do things with line...现在我转移到Python 3.8，它仍然可以正常处理纯文本文件，但是当它遇到gzipped文件时，我收到以下错误：TypeError: cannot use a string pattern on a bytes-like object解决这个问题的最佳方法是什么？我知道我可以检查是否line是一个字节对象并将其解码为字符串，但我宁愿使用一些标志来自动始终将行迭代为字符串（如果可能）；而且，我更喜欢编写同时适用于 Python 2 和 3 的代码。

Answer 1

fileinput.input根据是否获取 gzip 压缩文件，执行根本不同的操作。对于文本文件，它以常规打开open，默认情况下以文本模式有效打开。对于gzip.open，默认模式是二进制，这对于内容未知的压缩文件来说是合理的。

仅二进制限制是由 人为施加的fileinput.FileInput。从该方法的代码来看__init__：

  # restrict mode argument to reading modes

   if mode not in ('r', 'rU', 'U', 'rb'):

       raise ValueError("FileInput opening mode must be one of "

                        "'r', 'rU', 'U' and 'rb'")

   if 'U' in mode:

       import warnings

       warnings.warn("'U' mode is deprecated",

                     DeprecationWarning, 2)

   self._mode = mode

这为您提供了两种解决方法。

选项1

设置_mode后的属性__init__。为了避免在使用中添加额外的代码行，您可以子类化fileinput.FileInput并直接使用该类：

class TextFileInput(fileinput.FileInput):

    def __init__(*args, **kwargs):

        if 'mode' in kwargs and 't' in kwargs['mode']:

            mode = kwargs.pop['mode']

        else:

            mode = ''

        super().__init__(*args, **kwargs)

        if mode:

            self._mode = mode

for line in TextFileInput(filenames, openhook=fileinput.hook_compressed, mode='rt'):

    ...

选项2

弄乱未记录的前导下划线非常麻烦，因此您可以为 zip 文件创建自定义挂钩。这实际上非常简单，因为您可以使用代码作为fileinput.hook_compressed模板：

def my_hook_compressed(filename, mode):

    if 'b' not in mode:

        mode += 't'

    ext = os.path.splitext(filename)[1]

    if ext == '.gz':

        import gzip

        return gzip.open(filename, mode)

    elif ext == '.bz2':

        import bz2

        return bz2.open(filename, mode)

    else:

        return open(filename, mode)

选项3

最后，您始终可以将字节解码为 unicode 字符串。这显然不是更好的选择。

Answer 2

def my_hook_compressed(filename, mode):

    """hook for fileinput so we can also handle compressed files seamlessly"""

    if 'b' not in mode:

        mode += 't'

    ext = os.path.splitext(filename)[1]

    if ext == '.gz':

        import gzip

        return gzip.open(filename, mode)

    elif ext == '.bz2':

        import bz2

        return bz2.open(filename, mode)

    elif ext == '.xz':

        import lzma

        return lzma.open(filename, mode)

    elif ext == '.zst':

        import zstandard, io

        compressed = open(filename, 'rb')

        decompressor = zstandard.ZstdDecompressor()

        stream_reader = decompressor.stream_reader(compressed)

        return io.TextIOWrapper(stream_reader)

    else:

        return open(filename, mode)

我没有在 2.7 上测试过，但这适用于 3.8+

for line in fileinput.input(filenames, openhook=my_hook_compressed):

    ...