From da3676e94d22abf329cb492f7845ad8bf6918937 Mon Sep 17 00:00:00 2001
From: Mat Wood <matwood@meta.com>
Date: Tue, 5 Dec 2023 16:12:22 -0800
Subject: [PATCH] Add Source enum for capturing the source of the user-provided
 value

---
 Cargo.toml           |  2 +-
 src/file_or_stdin.rs | 26 +++++++++-----------------
 src/lib.rs           | 34 ++++++++++++++++++++++++++++++++++
 src/maybe_stdin.rs   | 26 +++++++++-----------------
 4 files changed, 53 insertions(+), 35 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c8c47fd..16f054c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "clap-stdin"
-version = "0.2.1"
+version = "0.2.2"
 edition = "2021"
 authors = ["Mat Wood <thepacketgeek@users.noreply.github.com>"]
 description = "Provides a type for easily accepting Clap arguments from stdin"
diff --git a/src/file_or_stdin.rs b/src/file_or_stdin.rs
index d389a28..bfca65d 100644
--- a/src/file_or_stdin.rs
+++ b/src/file_or_stdin.rs
@@ -1,9 +1,8 @@
 use std::fs;
 use std::io::{self, Read};
 use std::str::FromStr;
-use std::sync::atomic::Ordering;
 
-use super::{StdinError, STDIN_HAS_BEEN_USED};
+use super::{Source, StdinError};
 
 /// Wrapper struct to either read in a file or contents from `stdin`
 ///
@@ -34,15 +33,11 @@ use super::{StdinError, STDIN_HAS_BEEN_USED};
 /// ```
 #[derive(Clone)]
 pub struct FileOrStdin<T = String> {
+    /// Source of the contents
+    pub source: Source,
     inner: T,
 }
 
-impl<T> FileOrStdin<T> {
-    fn new(s: T) -> Self {
-        Self { inner: s }
-    }
-}
-
 impl<T> FromStr for FileOrStdin<T>
 where
     T: FromStr,
@@ -51,22 +46,19 @@ where
     type Err = StdinError;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "-" => {
-                if STDIN_HAS_BEEN_USED.load(std::sync::atomic::Ordering::Acquire) {
-                    return Err(StdinError::StdInRepeatedUse);
-                }
-                STDIN_HAS_BEEN_USED.store(true, Ordering::SeqCst);
+        let source = Source::from_str(s)?;
+        match &source {
+            Source::Stdin => {
                 let stdin = io::stdin();
                 let mut input = String::new();
                 stdin.lock().read_to_string(&mut input)?;
                 Ok(T::from_str(input.trim_end())
                     .map_err(|e| StdinError::FromStr(format!("{e}")))
-                    .map(|val| FileOrStdin::new(val))?)
+                    .map(|val| Self { source, inner: val })?)
             }
-            filepath => Ok(T::from_str(&fs::read_to_string(filepath)?)
+            Source::Arg(filepath) => Ok(T::from_str(&fs::read_to_string(filepath)?)
                 .map_err(|e| StdinError::FromStr(format!("{e}")))
-                .map(|val| FileOrStdin::new(val))?),
+                .map(|val| FileOrStdin { source, inner: val })?),
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 7b9ef62..035dbab 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,7 @@
 #![doc = include_str!("../README.md")]
 
 use std::io;
+use std::str::FromStr;
 use std::sync::atomic::AtomicBool;
 
 mod maybe_stdin;
@@ -19,3 +20,36 @@ pub enum StdinError {
     #[error("unable to parse from_str: {0}")]
     FromStr(String),
 }
+
+/// Source of the value contents will be either from `stdin` or a CLI arg provided value
+#[derive(Clone)]
+pub enum Source {
+    Stdin,
+    Arg(String),
+}
+
+impl FromStr for Source {
+    type Err = StdinError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "-" => {
+                if STDIN_HAS_BEEN_USED.load(std::sync::atomic::Ordering::Acquire) {
+                    return Err(StdinError::StdInRepeatedUse);
+                }
+                STDIN_HAS_BEEN_USED.store(true, std::sync::atomic::Ordering::SeqCst);
+                Ok(Self::Stdin)
+            }
+            arg => Ok(Self::Arg(arg.to_owned())),
+        }
+    }
+}
+
+impl std::fmt::Debug for Source {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Source::Stdin => write!(f, "stdin"),
+            Source::Arg(v) => v.fmt(f),
+        }
+    }
+}
diff --git a/src/maybe_stdin.rs b/src/maybe_stdin.rs
index 41706a0..712b247 100644
--- a/src/maybe_stdin.rs
+++ b/src/maybe_stdin.rs
@@ -1,8 +1,7 @@
 use std::io::{self, Read};
 use std::str::FromStr;
-use std::sync::atomic::Ordering;
 
-use super::{StdinError, STDIN_HAS_BEEN_USED};
+use super::{Source, StdinError};
 
 /// Wrapper struct to parse arg values from `stdin`
 ///
@@ -28,15 +27,11 @@ use super::{StdinError, STDIN_HAS_BEEN_USED};
 /// ```
 #[derive(Clone)]
 pub struct MaybeStdin<T> {
+    /// Source of the contents
+    pub source: Source,
     inner: T,
 }
 
-impl<T> MaybeStdin<T> {
-    fn new(inner: T) -> Self {
-        Self { inner }
-    }
-}
-
 impl<T> FromStr for MaybeStdin<T>
 where
     T: FromStr,
@@ -45,22 +40,19 @@ where
     type Err = StdinError;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "-" => {
-                if STDIN_HAS_BEEN_USED.load(std::sync::atomic::Ordering::Acquire) {
-                    return Err(StdinError::StdInRepeatedUse);
-                }
-                STDIN_HAS_BEEN_USED.store(true, Ordering::SeqCst);
+        let source = Source::from_str(s)?;
+        match &source {
+            Source::Stdin => {
                 let stdin = io::stdin();
                 let mut input = String::new();
                 stdin.lock().read_to_string(&mut input)?;
                 Ok(T::from_str(input.trim_end())
                     .map_err(|e| StdinError::FromStr(format!("{e}")))
-                    .map(|val| MaybeStdin::new(val))?)
+                    .map(|val| Self { source, inner: val })?)
             }
-            other => Ok(T::from_str(other)
+            Source::Arg(value) => Ok(T::from_str(value)
                 .map_err(|e| StdinError::FromStr(format!("{e}")))
-                .map(|val| MaybeStdin::new(val))?),
+                .map(|val| Self { source, inner: val })?),
         }
     }
 }