Add appropriate python exceptions (#8)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "cramjam"
-version = "1.1.0"
+version = "1.2.0"
 authors = ["Miles Granger <[email protected]>"]
 edition = "2018"
 license-file = "LICENSE"
@@ -14,6 +14,6 @@ crate-type = ["cdylib"]
 pyo3 = { version = "0.10.1", features = ["extension-module"] }
 snap = "^1"
 brotli2 = "^0.3"
-lz4-compress = "^0.1"
+lz4 = "1"
 flate2 = "^1"
 zstd = "0.5.1+zstd.1.4.4"

src/brotli.rs

@@ -1,18 +1,19 @@
 use brotli2::read::{BrotliDecoder, BrotliEncoder};
+use std::error::Error;
 use std::io::prelude::*;
 
 /// Decompress via Brotli
-pub fn decompress(data: &[u8]) -> Vec<u8> {
+pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut decoder = BrotliDecoder::new(data);
     let mut buf = vec![];
-    decoder.read_to_end(&mut buf).unwrap();
-    buf
+    decoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }
 
 /// Compress via Brotli
-pub fn compress(data: &[u8], level: u32) -> Vec<u8> {
+pub fn compress(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut encoder = BrotliEncoder::new(data, level);
     let mut buf = vec![];
-    encoder.read_to_end(&mut buf).unwrap();
-    buf
+    encoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }

src/deflate.rs

@@ -1,19 +1,20 @@
 use flate2::read::{DeflateDecoder, DeflateEncoder};
 use flate2::Compression;
+use std::error::Error;
 use std::io::prelude::*;
 
 /// Decompress gzip data
-pub fn decompress(data: &[u8]) -> Vec<u8> {
+pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut decoder = DeflateDecoder::new(data);
     let mut buf = vec![];
-    decoder.read_to_end(&mut buf).unwrap();
-    buf
+    decoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }
 
 /// Compress gzip data
-pub fn compress(data: &[u8], level: u32) -> Vec<u8> {
+pub fn compress(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut buf = vec![];
     let mut encoder = DeflateEncoder::new(data, Compression::new(level));
-    encoder.read_to_end(&mut buf).unwrap();
-    buf
+    encoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }

src/exceptions.rs

@@ -0,0 +1,5 @@
+use pyo3::create_exception;
+use pyo3::exceptions::Exception;
+
+create_exception!(cramjam, CompressionError, Exception);
+create_exception!(cramjam, DecompressionError, Exception);

src/gzip.rs

@@ -1,19 +1,20 @@
 use flate2::read::{GzDecoder, GzEncoder};
 use flate2::Compression;
+use std::error::Error;
 use std::io::prelude::*;
 
 /// Decompress gzip data
-pub fn decompress(data: &[u8]) -> Vec<u8> {
+pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut decoder = GzDecoder::new(data);
     let mut buf = vec![];
-    decoder.read_to_end(&mut buf).unwrap();
-    buf
+    decoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }
 
 /// Compress gzip data
-pub fn compress(data: &[u8], level: u32) -> Vec<u8> {
+pub fn compress(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn Error>> {
     let mut buf = vec![];
     let mut encoder = GzEncoder::new(data, Compression::new(level));
-    encoder.read_to_end(&mut buf).unwrap();
-    buf
+    encoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }

src/lib.rs

@@ -13,6 +13,7 @@
 
 mod brotli;
 mod deflate;
+mod exceptions;
 mod gzip;
 mod lz4;
 mod snappy;
@@ -22,6 +23,14 @@ use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use pyo3::wrap_pyfunction;
 
+use exceptions::{CompressionError, DecompressionError};
+
+macro_rules! to_py_err {
+    ($error:ident -> $expr:expr) => {
+        $expr.map_err(|err| PyErr::new::<$error, _>(err.to_string()))
+    };
+}
+
 /// Snappy decompression.
 ///
 /// Python Example
@@ -31,7 +40,7 @@ use pyo3::wrap_pyfunction;
 /// ```
 #[pyfunction]
 pub fn snappy_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = snappy::decompress(data);
+    let decompressed = to_py_err!(DecompressionError -> snappy::decompress(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -44,7 +53,7 @@ pub fn snappy_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyB
 /// ```
 #[pyfunction]
 pub fn snappy_compress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let compressed = snappy::compress(data);
+    let compressed = to_py_err!(CompressionError -> snappy::compress(data))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
@@ -58,7 +67,7 @@ pub fn snappy_compress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyByt
 /// ```
 #[pyfunction]
 pub fn snappy_decompress_raw<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = snappy::decompress_raw(data);
+    let decompressed = to_py_err!(DecompressionError -> snappy::decompress_raw(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -72,7 +81,7 @@ pub fn snappy_decompress_raw<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a
 /// ```
 #[pyfunction]
 pub fn snappy_compress_raw<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let compressed = snappy::compress_raw(data);
+    let compressed = to_py_err!(CompressionError -> snappy::compress_raw(data))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
@@ -85,7 +94,7 @@ pub fn snappy_compress_raw<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a P
 /// ```
 #[pyfunction]
 pub fn brotli_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = brotli::decompress(data);
+    let decompressed = to_py_err!(DecompressionError -> brotli::decompress(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -99,7 +108,7 @@ pub fn brotli_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyB
 #[pyfunction]
 pub fn brotli_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>) -> PyResult<&'a PyBytes> {
     let level = level.unwrap_or_else(|| 11);
-    let compressed = brotli::compress(data, level);
+    let compressed = to_py_err!(CompressionError -> brotli::compress(data, level))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
@@ -112,7 +121,7 @@ pub fn brotli_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>) -
 /// ```
 #[pyfunction]
 pub fn lz4_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = lz4::decompress(data);
+    let decompressed = to_py_err!(DecompressionError -> lz4::decompress(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -124,8 +133,9 @@ pub fn lz4_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyByte
 /// >>> lz4_compress(b'some bytes here')
 /// ```
 #[pyfunction]
-pub fn lz4_compress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let compressed = lz4::compress(data);
+pub fn lz4_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>) -> PyResult<&'a PyBytes> {
+    let level = level.unwrap_or_else(|| 4);
+    let compressed = to_py_err!(CompressionError -> lz4::compress(data, level))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
@@ -138,7 +148,7 @@ pub fn lz4_compress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes>
 /// ```
 #[pyfunction]
 pub fn gzip_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = gzip::decompress(data);
+    let decompressed = to_py_err!(DecompressionError -> gzip::decompress(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -152,7 +162,7 @@ pub fn gzip_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyByt
 #[pyfunction]
 pub fn gzip_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>) -> PyResult<&'a PyBytes> {
     let level = level.unwrap_or_else(|| 6);
-    let compressed = gzip::compress(data, level);
+    let compressed = to_py_err!(CompressionError -> gzip::compress(data, level))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
@@ -165,7 +175,7 @@ pub fn gzip_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>) ->
 /// ```
 #[pyfunction]
 pub fn deflate_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = deflate::decompress(data);
+    let decompressed = to_py_err!(DecompressionError -> deflate::decompress(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -179,7 +189,7 @@ pub fn deflate_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a Py
 #[pyfunction]
 pub fn deflate_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>) -> PyResult<&'a PyBytes> {
     let level = level.unwrap_or_else(|| 6);
-    let compressed = deflate::compress(data, level);
+    let compressed = to_py_err!(CompressionError -> deflate::compress(data, level))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
@@ -192,7 +202,7 @@ pub fn deflate_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<u32>)
 /// ```
 #[pyfunction]
 pub fn zstd_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyBytes> {
-    let decompressed = zstd::decompress(data);
+    let decompressed = to_py_err!(DecompressionError -> zstd::decompress(data))?;
     Ok(PyBytes::new(py, &decompressed))
 }
 
@@ -206,12 +216,15 @@ pub fn zstd_decompress<'a>(py: Python<'a>, data: &'a [u8]) -> PyResult<&'a PyByt
 #[pyfunction]
 pub fn zstd_compress<'a>(py: Python<'a>, data: &'a [u8], level: Option<i32>) -> PyResult<&'a PyBytes> {
     let level = level.unwrap_or_else(|| 0); // 0 will use zstd's default, currently 11
-    let compressed = zstd::compress(data, level);
+    let compressed = to_py_err!(CompressionError -> zstd::compress(data, level))?;
     Ok(PyBytes::new(py, &compressed))
 }
 
 #[pymodule]
-fn cramjam(_py: Python, m: &PyModule) -> PyResult<()> {
+fn cramjam(py: Python, m: &PyModule) -> PyResult<()> {
+    m.add("CompressionError", py.get_type::<CompressionError>())?;
+    m.add("DecompressionError", py.get_type::<DecompressionError>())?;
+
     m.add_wrapped(wrap_pyfunction!(snappy_compress))?;
     m.add_wrapped(wrap_pyfunction!(snappy_decompress))?;
     m.add_wrapped(wrap_pyfunction!(snappy_compress_raw))?;

src/lz4.rs

@@ -1,9 +1,22 @@
+use std::error::Error;
+use std::io::{Read, Write};
+
 /// Decompress lz4 data
-pub fn decompress(data: &[u8]) -> Vec<u8> {
-    lz4_compress::decompress(data).unwrap()
+pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Box<dyn Error>> {
+    let mut decoder = lz4::Decoder::new(data)?;
+    let mut buf = vec![];
+    decoder.read_to_end(&mut buf)?;
+    let (_, result) = decoder.finish(); // Weird API...
+    result?;
+    Ok(buf)
 }
 
 /// Compress lz4 data
-pub fn compress(data: &[u8]) -> Vec<u8> {
-    lz4_compress::compress(data)
+pub fn compress(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn Error>> {
+    let mut buf = vec![];
+    let mut encoder = lz4::EncoderBuilder::new().level(level).build(&mut buf)?;
+    encoder.write_all(data)?;
+    let (_, result) = encoder.finish(); // Weird API...
+    result?;
+    Ok(buf)
 }

src/snappy.rs

@@ -1,31 +1,31 @@
 use snap::raw::{Decoder, Encoder};
 use snap::read::{FrameDecoder, FrameEncoder};
-use std::io::Read;
+use std::io::{Error, Read};
 
 /// Decompress snappy data raw
-pub fn decompress_raw(data: &[u8]) -> Vec<u8> {
+pub fn decompress_raw(data: &[u8]) -> Result<Vec<u8>, snap::Error> {
     let mut decoder = Decoder::new();
-    decoder.decompress_vec(data).unwrap()
+    decoder.decompress_vec(data)
 }
 
 /// Compress snappy data raw
-pub fn compress_raw(data: &[u8]) -> Vec<u8> {
+pub fn compress_raw(data: &[u8]) -> Result<Vec<u8>, snap::Error> {
     let mut encoder = Encoder::new();
-    encoder.compress_vec(data).unwrap()
+    encoder.compress_vec(data)
 }
 
 /// Decompress snappy data framed
-pub fn decompress(data: &[u8]) -> Vec<u8> {
+pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Error> {
     let mut buf = vec![];
     let mut decoder = FrameDecoder::new(data);
-    decoder.read_to_end(&mut buf).unwrap();
-    buf
+    decoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }
 
 /// Decompress snappy data framed
-pub fn compress(data: &[u8]) -> Vec<u8> {
+pub fn compress(data: &[u8]) -> Result<Vec<u8>, Error> {
     let mut buf = vec![];
     let mut encoder = FrameEncoder::new(data);
-    encoder.read_to_end(&mut buf).unwrap();
-    buf
+    encoder.read_to_end(&mut buf)?;
+    Ok(buf)
 }

src/zstd.rs

@@ -1,9 +1,11 @@
+use std::io::Error;
+
 /// Decompress gzip data
-pub fn decompress(data: &[u8]) -> Vec<u8> {
-    zstd::stream::decode_all(data).unwrap()
+pub fn decompress(data: &[u8]) -> Result<Vec<u8>, Error> {
+    zstd::stream::decode_all(data)
 }
 
 /// Compress gzip data
-pub fn compress(data: &[u8], level: i32) -> Vec<u8> {
-    zstd::stream::encode_all(data, level).unwrap()
+pub fn compress(data: &[u8], level: i32) -> Result<Vec<u8>, Error> {
+    zstd::stream::encode_all(data, level)
 }

tests/test_variants.py

@@ -1,6 +1,7 @@
 import pytest
 
 import cramjam
+from cramjam import DecompressionError, CompressionError
 
 
 @pytest.mark.parametrize(
@@ -18,3 +19,11 @@ def test_variants_simple(variant):
 
     decompressed = decompress(compressed)
     assert decompressed == uncompressed
+
+
+@pytest.mark.parametrize(
+    "variant", ("snappy", "brotli", "lz4", "gzip", "deflate", "zstd")
+)
+def test_variants_raise_exception(variant):
+    with pytest.raises(cramjam.DecompressionError):
+        getattr(cramjam, f"{variant}_decompress")(b'sknow')