From d0aff43445f7d4827c4d778f1a41755c5e2d5155 Mon Sep 17 00:00:00 2001
From: Gregory Szorc <gregory.szorc@gmail.com>
Date: Wed, 30 Dec 2020 10:21:10 -0700
Subject: [PATCH] rust: start to implement ZstdCompressor.stream_writer()

I started to implemented this type then quickly found myself blocked
on https://github.com/PyO3/pyo3/issues/1205 /
https://github.com/PyO3/pyo3/issues/1206 due to not being able to return
`Self` from `__enter__`. We'll have to wait for a future pyo3 release
before we can finish the Rust port.
---
 rust-ext/src/compression_writer.rs | 46 ++++++++++++++++++++++++++++++
 rust-ext/src/compressor.rs         | 39 +++++++++++++++++++++++++
 rust-ext/src/lib.rs                |  1 +
 3 files changed, 86 insertions(+)
 create mode 100644 rust-ext/src/compression_writer.rs

diff --git a/rust-ext/src/compression_writer.rs b/rust-ext/src/compression_writer.rs
new file mode 100644
index 00000000..752cd697
--- /dev/null
+++ b/rust-ext/src/compression_writer.rs
@@ -0,0 +1,46 @@
+// Copyright (c) 2020-present, Gregory Szorc
+// All rights reserved.
+//
+// This software may be modified and distributed under the terms
+// of the BSD license. See the LICENSE file for details.
+
+use {crate::compressor::CCtx, pyo3::prelude::*, std::sync::Arc};
+
+#[pyclass]
+pub struct ZstdCompressionWriter {
+    cctx: Arc<CCtx<'static>>,
+    writer: PyObject,
+    source_size: u64,
+    write_size: usize,
+    write_return_read: bool,
+    closefd: bool,
+    entered: bool,
+    closing: bool,
+    closed: bool,
+    bytes_compressed: usize,
+}
+
+impl ZstdCompressionWriter {
+    pub fn new(
+        py: Python,
+        cctx: Arc<CCtx<'static>>,
+        writer: &PyAny,
+        source_size: u64,
+        write_size: usize,
+        write_return_read: bool,
+        closefd: bool,
+    ) -> Self {
+        Self {
+            cctx,
+            writer: writer.into_py(py),
+            source_size,
+            write_size,
+            write_return_read,
+            closefd,
+            entered: false,
+            closing: false,
+            closed: false,
+            bytes_compressed: 0,
+        }
+    }
+}
diff --git a/rust-ext/src/compressor.rs b/rust-ext/src/compressor.rs
index 24db16ef..50bc7c65 100644
--- a/rust-ext/src/compressor.rs
+++ b/rust-ext/src/compressor.rs
@@ -8,6 +8,7 @@ use {
     crate::{
         compression_dict::ZstdCompressionDict,
         compression_parameters::{CCtxParams, ZstdCompressionParameters},
+        compression_writer::ZstdCompressionWriter,
         compressionobj::ZstdCompressionObj,
         ZstdError,
     },
@@ -453,6 +454,44 @@ impl ZstdCompressor {
 
         Ok((total_read, total_write))
     }
+
+    #[args(
+        writer,
+        size = "None",
+        write_size = "None",
+        write_return_read = "true",
+        closefd = "true"
+    )]
+    fn stream_writer(
+        &self,
+        py: Python,
+        writer: &PyAny,
+        size: Option<u64>,
+        write_size: Option<usize>,
+        write_return_read: bool,
+        closefd: bool,
+    ) -> PyResult<ZstdCompressionWriter> {
+        if !writer.hasattr("write")? {
+            return Err(PyValueError::new_err(
+                "must pass object with a write() method",
+            ));
+        }
+
+        self.cctx.reset();
+
+        let size = size.unwrap_or(zstd_sys::ZSTD_CONTENTSIZE_UNKNOWN as _);
+        let write_size = write_size.unwrap_or_else(|| unsafe { zstd_sys::ZSTD_CStreamOutSize() });
+
+        Ok(ZstdCompressionWriter::new(
+            py,
+            self.cctx.clone(),
+            writer,
+            size,
+            write_size,
+            write_return_read,
+            closefd,
+        ))
+    }
 }
 
 pub(crate) fn init_module(module: &PyModule) -> PyResult<()> {
diff --git a/rust-ext/src/lib.rs b/rust-ext/src/lib.rs
index 6417249b..5eb47159 100644
--- a/rust-ext/src/lib.rs
+++ b/rust-ext/src/lib.rs
@@ -8,6 +8,7 @@ use pyo3::{prelude::*, types::PySet};
 
 mod compression_dict;
 mod compression_parameters;
+mod compression_writer;
 mod compressionobj;
 mod compressor;
 mod constants;