Using zstandard to compress a file in Python

Question

So I'm using the zstandard python library, and I've written a helper class and function to use contexts to decompress files.

class ZstdReader:
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        self.f = open(self.filename, 'rb')
        dctx = zstd.ZstdDecompressor()
        reader = dctx.stream_reader(self.f)
        return io.TextIOWrapper(reader, encoding='utf-8')

    def __exit__(self, *a):
        self.f.close()
        return False

def openZstd(filename, mode='rb'):
    if 'w' in mode:
        return ZstdWriter(filename)
    return ZstdReader(filename)

This works really well and allows me to just use with openZstd('filename.zst', 'rb') as f: before using the file f for json dumping and loading. I'm however having issues generalizing this to writing, I've tried following the documentation in the same way I did so far but something is not working. Here's what I've tried:

class ZstdWriter:
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        self.f = open(self.filename, 'wb')
        ctx = zstd.ZstdCompressor()
        writer = ctx.stream_writer(self.f)
        return io.TextIOWrapper(writer, encoding='utf-8')

    def __exit__(self, *a):
        self.f.close()
        return False

When I open a file using this class, and do a json.dump([], f), the file ends up being empty for some reason. I guess one of the steps is swallowing my input, but have no idea what it could possibly be.

I'm not familiar with `zstandard`, but I would imagine that either the `ctx` or `writer` objects need to have an explicit `close()` or `flush()` operation of some sort performed on them, to make sure that any buffered data actually gets written out. — jasonharper, Oct 22 '20 at 18:51
@jasonharper if I flush the wrapper, the writer and the file descriptor, then the file is non-empty and decompresses the right content but the `zstd` CLI complains about a premature end stil... — John Do, Oct 22 '20 at 21:53
From the docs: "If `flush(FLUSH_FRAME)` is not called, emitted data doesn’t constitute a full zstd frame and consumers of this data may complain about malformed input." Were you perhaps calling `flush()` with a different parameter? — jasonharper, Oct 22 '20 at 22:09

score 1 · Accepted Answer · answered Oct 24 '20 at 14:40

As suggested by jasonharper in the comments, you have to flush both the io wrapper and the writer itself, as follows:

s = json.dumps({})
iw = io.TextIOWrapper(writer, encoding="utf-8")
iw.write(s)

iw.flush()
writer.flush(zstd.FLUSH_FRAME)
f.close()

This results on the data being in the file, and the file being complete.

Using zstandard to compress a file in Python

1 Answers1