1 files changed, 37 insertions, 8 deletions

diff --git a/src/lib.rs b/src/lib.rs
index e5118be..23b4593 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -21,12 +21,12 @@
 //! Classic two-colored Hitomezashi stitch pattern generator.
 
 use rand::distributions::{Bernoulli, Distribution};
-use std::io::{stdout, BufWriter, StdoutLock, Write};
+use std::io::{stdout, BufWriter, Write};
 
 const TRANSPARENT_SQUARE: &[u8] = " ".as_bytes();
 const OPAQUE_SQUARE: &[u8] = "█".as_bytes();
 
-fn print_square(stream: &mut BufWriter<StdoutLock>, is_opaque: bool) {
+fn print_square<W: Write>(stream: &mut BufWriter<W>, is_opaque: bool) {
     let square = if is_opaque {
         OPAQUE_SQUARE
     } else {
@@ -36,7 +36,7 @@ fn print_square(stream: &mut BufWriter<StdoutLock>, is_opaque: bool) {
     stream.write_all(square).unwrap();
 }
 
-/// Prints a two-colored Hitomezashi stitch pattern of the specified dimensions.
+/// Writes a two-colored Hitomezashi stitch pattern of the specified dimensions to a buffer.
 ///
 /// `skew` is the probability of a row or column beginning with a stitch. Skew values near 0 or 1
 /// generate orderly patterns. Skew values near 0.5 generate chaotic patterns.
@@ -56,9 +56,14 @@ fn print_square(stream: &mut BufWriter<StdoutLock>, is_opaque: bool) {
 /// ```
 /// use hitomezashi_rs;
 ///
-/// hitomezashi_rs::generate(15, 20, Some(0.7));
+/// hitomezashi_rs::generate(15, 20, Some(0.7), &mut stream);
 /// ```
-pub fn generate(width: usize, height: usize, skew: Option<f64>) {
+fn generate_to_buffer<W: Write>(
+    width: usize,
+    height: usize,
+    skew: Option<f64>,
+    stream: &mut BufWriter<W>,
+) {
     // skew=0.5 generates the most random-looking patterns
     let skew = skew.unwrap_or(0.5);
 
@@ -102,9 +107,6 @@ pub fn generate(width: usize, height: usize, skew: Option<f64>) {
 
     let mut cur_row: Vec<bool> = Vec::with_capacity(width);
 
-    let lock = stdout().lock();
-    let mut stream = BufWriter::new(lock);
-
     /* Base case
      *
      * The first row has no preceding row, so it's computing using the column stitches */
@@ -134,3 +136,30 @@ pub fn generate(width: usize, height: usize, skew: Option<f64>) {
 
     stream.flush().unwrap();
 }
+
+/// Prints a two-colored Hitomezashi stitch pattern of the specified dimensions.
+///
+/// `skew` is the probability of a row or column beginning with a stitch. Skew values near 0 or 1
+/// generate orderly patterns. Skew values near 0.5 generate chaotic patterns.
+///
+/// # Panics
+///
+/// This function will panic if any of the following constraints are not met:
+///
+/// `width >= 1`
+///
+/// `height >= 1`
+///
+/// `0 <= skew <= 1`
+///
+/// # Examples
+///
+/// ```
+/// use hitomezashi_rs;
+///
+/// hitomezashi_rs::generate(15, 20, Some(0.7));
+/// ```
+pub fn generate(width: usize, height: usize, skew: Option<f64>) {
+    let mut stdout = BufWriter::new(stdout().lock());
+    generate_to_buffer(width, height, skew, &mut stdout);
+}