From a805d7cb1c8594aa2695968b164885628030f19c Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sun, 3 Jan 2021 13:05:35 +0800
Subject: [PATCH] feat: add stdout_or_file_write

---
 src/util_io.rs | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/util_io.rs b/src/util_io.rs
index 4e63ce6..e4ec2a1 100644
--- a/src/util_io.rs
+++ b/src/util_io.rs
@@ -1,13 +1,31 @@
-use std::{ fs::File, io::{ self, ErrorKind, prelude::* } };
+use std::io;
+use std::io::Write;
+use std::{fs::File, io::{ ErrorKind, prelude::* }};
 use std::time::{ SystemTime, Duration };
 
-use crate::{ XResult, new_box_ioerror };
+use crate::{SimpleError, XResult, new_box_ioerror};
 use crate::util_size;
 use crate::util_msg;
 use crate::util_file;
 
 pub const DEFAULT_BUF_SIZE: usize = 8 * 1024;
 
+pub fn stdout_or_file_write(file: Option<&str>, overwrite: bool) -> XResult<Box<dyn Write>> {
+    match file {
+        None => Ok(Box::new(io::stdout())),
+        Some(output) => {
+            if let Ok(_) = File::open(output) {
+                if !overwrite {
+                    return Err(SimpleError::new(format!("File exists: {}", output)).into());
+                }
+            }
+            Ok(Box::new(File::create(output).map_err(|e| {
+                SimpleError::new(format!("Create file: {}, failed: {}", output, e))
+            })?))
+        },
+    }
+}
+
 pub struct PrintStatusContext {
     pub print_interval_time: Duration,
     pub print_interval_bytes: i64,