Python zipfile + os.pipe()探索记

最新推荐文章于 2024-04-02 22:33:14 发布

weixin_34370347

最新推荐文章于 2024-04-02 22:33:14 发布

阅读量202

点赞数

文章标签： python java

原文链接：https://my.oschina.net/guyongquan/blog/1621632

版权

2019独角兽企业重金招聘Python工程师标准>>>

背景：需要使用Python把多个文件存为zip并通过HTTP上传压缩后的zip文件，不需要压缩，直接zip存储就可以了。

这能有多难，zipfile+requests就搞定了。
唯一的小问题是，先压缩成一个本地文件然后上传显然不太优雅，用IO管道流能一边存储为zip一边上传最好了。
想来也不是什么大问题，用Java的话分分钟搞定，只是Python不太熟。

花上几分钟看看Python的os.pipe()，代码很快撸出来了。执行，文件上传成功。只是，解压报错说zip文件不合法。

为啥了？
上传过程中出错？改改代码不上传数据，直接保存为本地文件，解压依然报错。上传过程中出错的可能性排除。
不通过os.pipe直接压缩到本地文件，解压成功。加上os.pipe中转一下就报错。看来就是这个环节错误了。

对比一下错误的文件和正确的文件，发现文件头部就存在明显差异。

瞄一下zipfile源代码，下面两行行引起了注意：

            self._fileobj.seek(self._zinfo.header_offset)
            self._fileobj.write(self._zinfo.FileHeader(self._zip64))

seek然后写入文件头信息，可是我这是管道流啊，肯定没法seek回去修改数据的。

为什么要seek回去写文件头了？而且之前用java做过类似的事情，没遇到这问题。
带着疑问去看zip文件格式，原来是文件头信息里的crc校验码和压缩后大小需要seek回去写。不过也说了这两个东西是可以写到后面去的。
但是zipfile并没有提供这样的选择设置。

怎么办了？
把zipfile源代码复制一下改改？太粗暴了。
先看看能不能在传给zipfile的io流对象上做手脚，看能不能重写seek和write方法，然后修改zip头里的标志位改了，把crc校验码写到文件后面去。

那先看看压缩过程中有多少地方调用了seek。在Python标准库的IO流对象seek方法下断点，死活不命中。可能是因为这些标准库实现里使用了大量高级Python特性，导致类定义里的方法其实根本只是一个签名，真正的实现根本不在类定义里。Python这点真是，挺讨厌的。

那自己定义一个IO流对象，手写seek方法下断点好了。
果然命中，还有惊喜。
原来zipfile里会判断流对象是否能seek，不能的话就会把crc校验码那些东西写到后面去，这样就不用seek了。

这就简单啦，写个wrapper包装一下pipe流，直接在seek方法被调用的时候抛出一个AttributeError就好了。利用Python那些有点讨厌的高级特性做这事倒是比java简单多了，重写__getattribute__就可以了。

事实证明我还是想太简单，文件头是对了。文件结尾不对。
继续对照zip文件格式看不对的字节，错误的字节是Central Directory的开始位置偏移。

既然是文件末尾不对，那应该是close的时候写入的数据错误，直觉应该是zipfile里计算字节数错误。
在zipfile源代码的close方法下断点，启动调试，然后想起zipfile里判断io流是否能seek的代码附近还会判断io流是否能tell的。
不会是管道流的tell方法实现也不一样吧，直接tell方法也重写抛出AttributeError得了。
嘿，还真可以了。

import os
from io import BufferedIOBase
import zipfile


class _DisableSeekAndTellIOWrapper(BufferedIOBase):
    """
    zipfile压缩文件的时候会通过seek和tell来定位写入zip文件头信息。
    但是管道流不能正常的seek和tell，故通过这个wrapper来禁用seek和tell，从而强迫zipfile使用流式压缩
    """

    def __init__(self, towrap):
        self._wrapped = towrap

    def seek(self, *args, **kwargs):
        raise AttributeError()

    def tell(self, *args, **kwargs):
        raise AttributeError()

    def seekable(self, *args, **kwargs):
        return False

    def __getattribute__(self, item):
        if item in ('_wrapped', 'seek', 'seekable', 'tell'):
            return object.__getattribute__(self, item)
        return self._wrapped.__getattribute__(item)

    def __enter__(self):
        self._wrapped.__enter__()
        return self

    def __exit__(self, *args, **kwargs):
        self._wrapped.__exit__(*args, **kwargs)


def _zip_directory_to_pipe(dir_path: str) -> BufferedIOBase:
    rfd, wfd = os.pipe()

    def zip_in_thread(wfd):
        with os.fdopen(wfd, 'wb') as outputfile:
            with zipfile.ZipFile(_DisableSeekAndTellIOWrapper(outputfile), 'w') as zfile:
                for root, dirs, files in os.walk(dir_path):
                    for name in files:
                        fullpath = os.path.join(root, name)
                        arcname = os.path.relpath(fullpath, dir_path)
                        zfile.write(fullpath, arcname=arcname)

    import threading
    threading.Thread(target=zip_in_thread, args=(wfd,), name='zip ' + dir_path).start()
    return os.fdopen(rfd, 'rb')

import requests

requests.post('http://localhost:8080/upload', files = {'file': _zip_directory_to_pipe('d:/testdir')})

转载于:https://my.oschina.net/guyongquan/blog/1621632