From 08f0e73926c11adc3105dbf4eb84dd8c9e6c873a Mon Sep 17 00:00:00 2001
From: Andrew Chin <achin@eminence32.net>
Date: Thu, 31 Aug 2023 23:58:39 -0400
Subject: [PATCH] Check all streams (out/err/in) to find one that has a tty and
 size

Closes #48
Closes #23
---
 src/lib.rs     |  4 +++-
 src/unix.rs    | 18 ++++++++++++++----
 src/windows.rs | 28 +++++++++++++++++++++++-----
 3 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index e933b21..6f8990c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -34,4 +34,6 @@ mod windows;
 pub use crate::windows::{terminal_size, terminal_size_using_handle};
 
 #[cfg(not(any(unix, windows)))]
-pub fn terminal_size() -> Option<(Width, Height)> { None }
+pub fn terminal_size() -> Option<(Width, Height)> {
+    None
+}
diff --git a/src/unix.rs b/src/unix.rs
index 59af979..b218b1e 100644
--- a/src/unix.rs
+++ b/src/unix.rs
@@ -1,12 +1,22 @@
 use super::{Height, Width};
-use std::os::unix::io::RawFd;
 use rustix::fd::BorrowedFd;
+use std::os::unix::io::RawFd;
 
-/// Returns the size of the terminal defaulting to STDOUT, if available.
+/// Returns the size of the terminal.
 ///
-/// If STDOUT is not a tty, returns `None`
+/// This function checks the stdout, stderr, and stdin streams (in that order).
+/// The size of the first stream that is a TTY will be returned.  If nothing
+/// is a TTY, then `None` is returned.
 pub fn terminal_size() -> Option<(Width, Height)> {
-    terminal_size_using_fd(rustix::io::raw_stdout())
+    if let Some(size) = terminal_size_using_fd(rustix::io::raw_stdout()) {
+        Some(size)
+    } else if let Some(size) = terminal_size_using_fd(rustix::io::raw_stderr()) {
+        Some(size)
+    } else if let Some(size) = terminal_size_using_fd(rustix::io::raw_stdin()) {
+        Some(size)
+    } else {
+        None
+    }
 }
 
 /// Returns the size of the terminal using the given file descriptor, if available.
diff --git a/src/windows.rs b/src/windows.rs
index 74315d6..ff98b86 100644
--- a/src/windows.rs
+++ b/src/windows.rs
@@ -1,16 +1,34 @@
 use super::{Height, Width};
 use std::os::windows::io::RawHandle;
 
-/// Returns the size of the terminal defaulting to STDOUT, if available.
+/// Returns the size of the terminal.
+///
+/// This function checks the stdout, stderr, and stdin streams (in that order).
+/// The size of the first stream that is a TTY will be returned.  If nothing
+/// is a TTY, then `None` is returned.
 ///
 /// Note that this returns the size of the actual command window, and
 /// not the overall size of the command window buffer
 pub fn terminal_size() -> Option<(Width, Height)> {
-    use windows_sys::Win32::System::Console::{GetStdHandle, STD_OUTPUT_HANDLE};
-
-    let handle = unsafe { GetStdHandle(STD_OUTPUT_HANDLE) as RawHandle };
+    use windows_sys::Win32::System::Console::{
+        GetStdHandle, STD_ERROR_HANDLE, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE,
+    };
 
-    terminal_size_using_handle(handle)
+    if let Some(size) =
+        terminal_size_using_handle(unsafe { GetStdHandle(STD_OUTPUT_HANDLE) as RawHandle })
+    {
+        Some(size)
+    } else if let Some(size) =
+        terminal_size_using_handle(unsafe { GetStdHandle(STD_ERROR_HANDLE) as RawHandle })
+    {
+        Some(size)
+    } else if let Some(size) =
+        terminal_size_using_handle(unsafe { GetStdHandle(STD_INPUT_HANDLE) as RawHandle })
+    {
+        Some(size)
+    } else {
+        None
+    }
 }
 
 /// Returns the size of the terminal using the given handle, if available.