--- a/contrib/python-zstandard/c-ext/compressionwriter.c	Thu Apr 04 15:24:03 2019 -0700
+++ b/contrib/python-zstandard/c-ext/compressionwriter.c	Thu Apr 04 17:34:43 2019 -0700
@@ -18,24 +18,23 @@
 	Py_XDECREF(self->compressor);
 	Py_XDECREF(self->writer);
 
+	PyMem_Free(self->output.dst);
+	self->output.dst = NULL;
+
 	PyObject_Del(self);
 }
 
 static PyObject* ZstdCompressionWriter_enter(ZstdCompressionWriter* self) {
-	size_t zresult;
+	if (self->closed) {
+		PyErr_SetString(PyExc_ValueError, "stream is closed");
+		return NULL;
+	}
 
 	if (self->entered) {
 		PyErr_SetString(ZstdError, "cannot __enter__ multiple times");
 		return NULL;
 	}
 
-	zresult = ZSTD_CCtx_setPledgedSrcSize(self->compressor->cctx, self->sourceSize);
-	if (ZSTD_isError(zresult)) {
-		PyErr_Format(ZstdError, "error setting source size: %s",
-			ZSTD_getErrorName(zresult));
-		return NULL;
-	}
-
 	self->entered = 1;
 
 	Py_INCREF(self);
@@ -46,10 +45,6 @@
 	PyObject* exc_type;
 	PyObject* exc_value;
 	PyObject* exc_tb;
-	size_t zresult;
-
-	ZSTD_outBuffer output;
-	PyObject* res;
 
 	if (!PyArg_ParseTuple(args, "OOO:__exit__", &exc_type, &exc_value, &exc_tb)) {
 		return NULL;
@@ -58,46 +53,11 @@
 	self->entered = 0;
 
 	if (exc_type == Py_None && exc_value == Py_None && exc_tb == Py_None) {
-		ZSTD_inBuffer inBuffer;
-
-		inBuffer.src = NULL;
-		inBuffer.size = 0;
-		inBuffer.pos = 0;
-
-		output.dst = PyMem_Malloc(self->outSize);
-		if (!output.dst) {
-			return PyErr_NoMemory();
-		}
-		output.size = self->outSize;
-		output.pos = 0;
+		PyObject* result = PyObject_CallMethod((PyObject*)self, "close", NULL);
 
-		while (1) {
-			zresult = ZSTD_compress_generic(self->compressor->cctx, &output, &inBuffer, ZSTD_e_end);
-			if (ZSTD_isError(zresult)) {
-				PyErr_Format(ZstdError, "error ending compression stream: %s",
-					ZSTD_getErrorName(zresult));
-				PyMem_Free(output.dst);
-				return NULL;
-			}
-
-			if (output.pos) {
-#if PY_MAJOR_VERSION >= 3
-				res = PyObject_CallMethod(self->writer, "write", "y#",
-#else
-				res = PyObject_CallMethod(self->writer, "write", "s#",
-#endif
-					output.dst, output.pos);
-				Py_XDECREF(res);
-			}
-
-			if (!zresult) {
-				break;
-			}
-
-			output.pos = 0;
+		if (NULL == result) {
+			return NULL;
 		}
-
-		PyMem_Free(output.dst);
 	}
 
 	Py_RETURN_FALSE;
@@ -117,7 +77,6 @@
 	Py_buffer source;
 	size_t zresult;
 	ZSTD_inBuffer input;
-	ZSTD_outBuffer output;
 	PyObject* res;
 	Py_ssize_t totalWrite = 0;
 
@@ -130,143 +89,240 @@
 		return NULL;
 	}
 
-	if (!self->entered) {
-		PyErr_SetString(ZstdError, "compress must be called from an active context manager");
-		goto finally;
-	}
-
 	if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) {
 		PyErr_SetString(PyExc_ValueError,
 			"data buffer should be contiguous and have at most one dimension");
 		goto finally;
 	}
 
-	output.dst = PyMem_Malloc(self->outSize);
-	if (!output.dst) {
-		PyErr_NoMemory();
-		goto finally;
+	if (self->closed) {
+		PyErr_SetString(PyExc_ValueError, "stream is closed");
+		return NULL;
 	}
-	output.size = self->outSize;
-	output.pos = 0;
+
+	self->output.pos = 0;
 
 	input.src = source.buf;
 	input.size = source.len;
 	input.pos = 0;
 
-	while ((ssize_t)input.pos < source.len) {
+	while (input.pos < (size_t)source.len) {
 		Py_BEGIN_ALLOW_THREADS
-		zresult = ZSTD_compress_generic(self->compressor->cctx, &output, &input, ZSTD_e_continue);
+		zresult = ZSTD_compressStream2(self->compressor->cctx, &self->output, &input, ZSTD_e_continue);
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
-			PyMem_Free(output.dst);
 			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
 			goto finally;
 		}
 
 		/* Copy data from output buffer to writer. */
-		if (output.pos) {
+		if (self->output.pos) {
 #if PY_MAJOR_VERSION >= 3
 			res = PyObject_CallMethod(self->writer, "write", "y#",
 #else
 			res = PyObject_CallMethod(self->writer, "write", "s#",
 #endif
-				output.dst, output.pos);
+				self->output.dst, self->output.pos);
 			Py_XDECREF(res);
-			totalWrite += output.pos;
-			self->bytesCompressed += output.pos;
+			totalWrite += self->output.pos;
+			self->bytesCompressed += self->output.pos;
 		}
-		output.pos = 0;
+		self->output.pos = 0;
 	}
 
-	PyMem_Free(output.dst);
-
-	result = PyLong_FromSsize_t(totalWrite);
+	if (self->writeReturnRead) {
+		result = PyLong_FromSize_t(input.pos);
+	}
+	else {
+		result = PyLong_FromSsize_t(totalWrite);
+	}
 
 finally:
 	PyBuffer_Release(&source);
 	return result;
 }
 
-static PyObject* ZstdCompressionWriter_flush(ZstdCompressionWriter* self, PyObject* args) {
+static PyObject* ZstdCompressionWriter_flush(ZstdCompressionWriter* self, PyObject* args, PyObject* kwargs) {
+	static char* kwlist[] = {
+		"flush_mode",
+		NULL
+	};
+
 	size_t zresult;
-	ZSTD_outBuffer output;
 	ZSTD_inBuffer input;
 	PyObject* res;
 	Py_ssize_t totalWrite = 0;
+	unsigned flush_mode = 0;
+	ZSTD_EndDirective flush;
 
-	if (!self->entered) {
-		PyErr_SetString(ZstdError, "flush must be called from an active context manager");
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|I:flush",
+		kwlist, &flush_mode)) {
 		return NULL;
 	}
 
+	switch (flush_mode) {
+		case 0:
+			flush = ZSTD_e_flush;
+			break;
+		case 1:
+			flush = ZSTD_e_end;
+			break;
+		default:
+			PyErr_Format(PyExc_ValueError, "unknown flush_mode: %d", flush_mode);
+			return NULL;
+	}
+
+	if (self->closed) {
+		PyErr_SetString(PyExc_ValueError, "stream is closed");
+		return NULL;
+	}
+
+	self->output.pos = 0;
+
 	input.src = NULL;
 	input.size = 0;
 	input.pos = 0;
 
-	output.dst = PyMem_Malloc(self->outSize);
-	if (!output.dst) {
-		return PyErr_NoMemory();
-	}
-	output.size = self->outSize;
-	output.pos = 0;
-
 	while (1) {
 		Py_BEGIN_ALLOW_THREADS
-		zresult = ZSTD_compress_generic(self->compressor->cctx, &output, &input, ZSTD_e_flush);
+		zresult = ZSTD_compressStream2(self->compressor->cctx, &self->output, &input, flush);
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
-			PyMem_Free(output.dst);
 			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
 			return NULL;
 		}
 
 		/* Copy data from output buffer to writer. */
-		if (output.pos) {
+		if (self->output.pos) {
 #if PY_MAJOR_VERSION >= 3
 			res = PyObject_CallMethod(self->writer, "write", "y#",
 #else
 			res = PyObject_CallMethod(self->writer, "write", "s#",
 #endif
-				output.dst, output.pos);
+				self->output.dst, self->output.pos);
 			Py_XDECREF(res);
-			totalWrite += output.pos;
-			self->bytesCompressed += output.pos;
+			totalWrite += self->output.pos;
+			self->bytesCompressed += self->output.pos;
 		}
 
-		output.pos = 0;
+		self->output.pos = 0;
 
 		if (!zresult) {
 			break;
 		}
 	}
 
-	PyMem_Free(output.dst);
+	return PyLong_FromSsize_t(totalWrite);
+}
+
+static PyObject* ZstdCompressionWriter_close(ZstdCompressionWriter* self) {
+	PyObject* result;
+
+	if (self->closed) {
+		Py_RETURN_NONE;
+	}
+
+	result = PyObject_CallMethod((PyObject*)self, "flush", "I", 1);
+	self->closed = 1;
+
+	if (NULL == result) {
+	    return NULL;
+	}
 
-	return PyLong_FromSsize_t(totalWrite);
+    /* Call close on underlying stream as well. */
+	if (PyObject_HasAttrString(self->writer, "close")) {
+		return PyObject_CallMethod(self->writer, "close", NULL);
+	}
+
+	Py_RETURN_NONE;
+}
+
+static PyObject* ZstdCompressionWriter_fileno(ZstdCompressionWriter* self) {
+	if (PyObject_HasAttrString(self->writer, "fileno")) {
+		return PyObject_CallMethod(self->writer, "fileno", NULL);
+	}
+	else {
+		PyErr_SetString(PyExc_OSError, "fileno not available on underlying writer");
+		return NULL;
+	}
 }
 
 static PyObject* ZstdCompressionWriter_tell(ZstdCompressionWriter* self) {
 	return PyLong_FromUnsignedLongLong(self->bytesCompressed);
 }
 
+static PyObject* ZstdCompressionWriter_writelines(PyObject* self, PyObject* args) {
+	PyErr_SetNone(PyExc_NotImplementedError);
+	return NULL;
+}
+
+static PyObject* ZstdCompressionWriter_false(PyObject* self, PyObject* args) {
+	Py_RETURN_FALSE;
+}
+
+static PyObject* ZstdCompressionWriter_true(PyObject* self, PyObject* args) {
+	Py_RETURN_TRUE;
+}
+
+static PyObject* ZstdCompressionWriter_unsupported(PyObject* self, PyObject* args, PyObject* kwargs) {
+	PyObject* iomod;
+	PyObject* exc;
+
+	iomod = PyImport_ImportModule("io");
+	if (NULL == iomod) {
+		return NULL;
+	}
+
+	exc = PyObject_GetAttrString(iomod, "UnsupportedOperation");
+	if (NULL == exc) {
+		Py_DECREF(iomod);
+		return NULL;
+	}
+
+	PyErr_SetNone(exc);
+	Py_DECREF(exc);
+	Py_DECREF(iomod);
+
+	return NULL;
+}
+
 static PyMethodDef ZstdCompressionWriter_methods[] = {
 	{ "__enter__", (PyCFunction)ZstdCompressionWriter_enter, METH_NOARGS,
 	PyDoc_STR("Enter a compression context.") },
 	{ "__exit__", (PyCFunction)ZstdCompressionWriter_exit, METH_VARARGS,
 	PyDoc_STR("Exit a compression context.") },
+	{ "close", (PyCFunction)ZstdCompressionWriter_close, METH_NOARGS, NULL },
+	{ "fileno", (PyCFunction)ZstdCompressionWriter_fileno, METH_NOARGS, NULL },
+	{ "isatty", (PyCFunction)ZstdCompressionWriter_false, METH_NOARGS, NULL },
+	{ "readable", (PyCFunction)ZstdCompressionWriter_false, METH_NOARGS, NULL },
+	{ "readline", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+	{ "readlines", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+	{ "seek", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+	{ "seekable", ZstdCompressionWriter_false, METH_NOARGS, NULL },
+	{ "truncate", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+	{ "writable", ZstdCompressionWriter_true, METH_NOARGS, NULL },
+	{ "writelines", ZstdCompressionWriter_writelines, METH_VARARGS, NULL },
+	{ "read", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+	{ "readall", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+	{ "readinto", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
 	{ "memory_size", (PyCFunction)ZstdCompressionWriter_memory_size, METH_NOARGS,
 	PyDoc_STR("Obtain the memory size of the underlying compressor") },
 	{ "write", (PyCFunction)ZstdCompressionWriter_write, METH_VARARGS | METH_KEYWORDS,
 	PyDoc_STR("Compress data") },
-	{ "flush", (PyCFunction)ZstdCompressionWriter_flush, METH_NOARGS,
+	{ "flush", (PyCFunction)ZstdCompressionWriter_flush, METH_VARARGS | METH_KEYWORDS,
 	PyDoc_STR("Flush data and finish a zstd frame") },
 	{ "tell", (PyCFunction)ZstdCompressionWriter_tell, METH_NOARGS,
 	PyDoc_STR("Returns current number of bytes compressed") },
 	{ NULL, NULL }
 };
 
+static PyMemberDef ZstdCompressionWriter_members[] = {
+	 { "closed", T_BOOL, offsetof(ZstdCompressionWriter, closed), READONLY, NULL },
+	 { NULL }
+};
+
 PyTypeObject ZstdCompressionWriterType = {
 	PyVarObject_HEAD_INIT(NULL, 0)
 	"zstd.ZstdCompressionWriter",  /* tp_name */
@@ -296,7 +352,7 @@
 	0,                              /* tp_iter */
 	0,                              /* tp_iternext */
 	ZstdCompressionWriter_methods,  /* tp_methods */
-	0,                              /* tp_members */
+	ZstdCompressionWriter_members,  /* tp_members */
 	0,                              /* tp_getset */
 	0,                              /* tp_base */
 	0,                              /* tp_dict */