[GR-13036] Implement zlib.decompressobj

Author: timfel

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/ZLibModuleBuiltins.java

@@ -453,6 +453,76 @@ Object deflateCompress(DeflaterWrapper stream, PIBytesLike pb, int mode) {
         }
     }
 
+    @Builtin(name = "zlib_inflateInit", fixedNumOfPositionalArgs = 2)
+    @GenerateNodeFactory
+    abstract static class InflateInitNode extends PythonBinaryBuiltinNode {
+        @Child BytesNodes.ToBytesNode toBytes = BytesNodes.ToBytesNode.create();
+
+        @Specialization
+        @TruffleBoundary
+        Object init(int wbits, PBytes zdict) {
+            Inflater inflater;
+            if (wbits < 0) {
+                // generate a RAW stream, i.e., no wrapping
+                inflater = new Inflater(true);
+            } else if (wbits >= 25) {
+                // include gzip container
+                throw raise(PythonBuiltinClassType.NotImplementedError, "gzip containers");
+            } else {
+                // wrap stream with zlib header and trailer
+                inflater = new Inflater(false);
+            }
+
+            inflater.setDictionary(toBytes.execute(zdict));
+            return new InflaterWrapper(inflater);
+        }
+    }
+
+    static class InflaterWrapper implements TruffleObject {
+        private final Inflater inflater;
+
+        public InflaterWrapper(Inflater inflater) {
+            this.inflater = inflater;
+        }
+
+        public ForeignAccess getForeignAccess() {
+            return null;
+        }
+    }
+
+    @Builtin(name = "zlib_inflateDecompress", fixedNumOfPositionalArgs = 3)
+    @GenerateNodeFactory
+    abstract static class InflaterDecompress extends PythonTernaryBuiltinNode {
+        @Child BytesNodes.ToBytesNode toBytes = BytesNodes.ToBytesNode.create();
+
+        @Specialization
+        @TruffleBoundary
+        Object decompress(InflaterWrapper stream, PIBytesLike pb, int maxLen) {
+            int maxLength = maxLen == 0 ? Integer.MAX_VALUE : maxLen;
+
+            ByteArrayOutputStream baos = new ByteArrayOutputStream();
+            byte[] data = toBytes.execute(pb);
+            byte[] result = new byte[DEF_BUF_SIZE];
+            stream.inflater.setInput(data);
+
+            int bytesWritten = result.length;
+            while (baos.size() < maxLength && bytesWritten == result.length) {
+                try {
+                    bytesWritten = stream.inflater.inflate(result, 0, result.length);
+                } catch (DataFormatException e) {
+                    throw raise(ZLibError, e.getMessage());
+                }
+                baos.write(result, 0, bytesWritten);
+            }
+
+            return factory().createTuple(new Object[]{
+                            factory().createBytes(baos.toByteArray()),
+                            stream.inflater.needsInput(),
+                            stream.inflater.getRemaining()
+            });
+        }
+    }
+
     // zlib.compress(data, level=-1)
     @Builtin(name = "compress", minNumOfPositionalArgs = 1, maxNumOfPositionalArgs = 2, keywordArguments = {"level"})
     @TypeSystemReference(PythonArithmeticTypes.class)

graalpython/lib-graalpython/zlib.py

@@ -104,3 +104,44 @@ def flush(self, mode=Z_FINISH):
             return result
         else:
             raise error("compressor object already flushed")
+
+
+class decompressobj():
+    """decompressobj([wbits]) -- Return a decompressor object.
+
+    Optional arg wbits is the window buffer size.
+    """
+    def __new__(cls, wbits=MAX_WBITS, zdict=b""):
+        self = object.__new__(cls)
+        self.unused_data = b""
+        self.unconsumed_tail = b""
+        self.eof = False
+        self.stream = zlib_inflateInit(wbits, zdict)
+        return self
+
+    def decompress(self, data, max_length=0):
+        """
+        decompress(data[, max_length]) -- Return a string containing the
+        decompressed version of the data.
+
+        If the max_length parameter is specified then the return value will be
+        no longer than max_length.  Unconsumed input data will be stored in the
+        unconsumed_tail attribute.
+        """
+        if max_length < 0:
+            raise ValueError("max_length must be greater than zero")
+        result, self.eof, unused_len = zlib_inflateDecompress(self.stream, data, max_length)
+        tail = data[len(data) - unused_len:]
+        if self.eof:
+            self.unconsumed_tail = b""
+            self.unused_data += tail
+        else:
+            self.unconsumed_tail = tail
+
+        return result
+
+    def flush(self, length=None):
+        try:
+            return decompress(self.unconsumed_tail)
+        except error:
+            return b""