From cd64884180ea82e02ea70e8dab2d59b735dcc6fe Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Thu, 9 Jul 2020 09:39:37 -0500
Subject: [PATCH] Implement os.get_terminal_size

---
 vm/Cargo.toml           |  3 ++-
 vm/src/stdlib/os.rs     | 57 +++++++++++++++++++++++++++++++++++++++++
 vm/src/stdlib/winapi.rs |  7 +++--
 3 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/vm/Cargo.toml b/vm/Cargo.toml
index ea66f69af..289fa2481 100644
--- a/vm/Cargo.toml
+++ b/vm/Cargo.toml
@@ -125,7 +125,8 @@ schannel = "0.1"
 version = "0.3"
 features = [
   "winsock2", "handleapi", "ws2def", "std", "winbase", "wincrypt", "fileapi", "processenv",
-  "namedpipeapi", "winnt", "processthreadsapi", "errhandlingapi", "winuser", "synchapi",
+  "namedpipeapi", "winnt", "processthreadsapi", "errhandlingapi", "winuser", "synchapi", "wincon",
+  "impl-default",
 ]
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
index 56cbe39cc..4eb33c3f2 100644
--- a/vm/src/stdlib/os.rs
+++ b/vm/src/stdlib/os.rs
@@ -1867,6 +1867,60 @@ fn os_strerror(e: i32) -> String {
         .into_owned()
 }
 
+#[pystruct_sequence(name = "os.terminal_size")]
+#[allow(dead_code)]
+struct PyTerminalSize {
+    columns: usize,
+    lines: usize,
+}
+
+fn os_get_terminal_size(fd: OptionalArg<i32>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
+    let (columns, lines) = {
+        #[cfg(unix)]
+        {
+            nix::ioctl_read_bad!(winsz, libc::TIOCGWINSZ, libc::winsize);
+            let mut w = libc::winsize {
+                ws_row: 0,
+                ws_col: 0,
+                ws_xpixel: 0,
+                ws_ypixel: 0,
+            };
+            unsafe { winsz(fd.unwrap_or(libc::STDOUT_FILENO), &mut w) }
+                .map_err(|e| convert_nix_error(vm, e))?;
+            (w.ws_col.into(), w.ws_row.into())
+        }
+        #[cfg(windows)]
+        {
+            use winapi::um::{handleapi, processenv, winbase, wincon};
+            let stdhandle = match fd {
+                OptionalArg::Present(0) => winbase::STD_INPUT_HANDLE,
+                OptionalArg::Present(1) | OptionalArg::Missing => winbase::STD_OUTPUT_HANDLE,
+                OptionalArg::Present(2) => winbase::STD_ERROR_HANDLE,
+                _ => return Err(vm.new_value_error("bad file descriptor".to_owned())),
+            };
+            let h = unsafe { processenv::GetStdHandle(stdhandle) };
+            if h.is_null() {
+                return Err(vm.new_os_error("handle cannot be retrieved".to_owned()));
+            }
+            if h == handleapi::INVALID_HANDLE_VALUE {
+                return Err(errno_err(vm));
+            }
+            let mut csbi = wincon::CONSOLE_SCREEN_BUFFER_INFO::default();
+            let ret = unsafe { wincon::GetConsoleScreenBufferInfo(h, &mut csbi) };
+            if ret == 0 {
+                return Err(errno_err(vm));
+            }
+            let w = csbi.srWindow;
+            (
+                (w.Right - w.Left + 1) as usize,
+                (w.Bottom - w.Top + 1) as usize,
+            )
+        }
+    };
+    PyTerminalSize { columns, lines }
+        .into_struct_sequence(vm, vm.try_class(MODULE_NAME, "terminal_size")?)
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
@@ -1885,6 +1939,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     });
 
     let stat_result = StatResult::make_class(ctx);
+    let terminal_size = PyTerminalSize::make_class(ctx);
 
     struct SupportFunc<'a> {
         name: &'a str,
@@ -1964,6 +2019,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "ScandirIter" => scandir_iter,
         "DirEntry" => dir_entry,
         "stat_result" => stat_result,
+        "terminal_size" => terminal_size,
         "lstat" => ctx.new_function(os_lstat),
         "getcwd" => ctx.new_function(os_getcwd),
         "chdir" => ctx.new_function(os_chdir),
@@ -1978,6 +2034,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "link" => ctx.new_function(os_link),
         "kill" => ctx.new_function(os_kill),
         "strerror" => ctx.new_function(os_strerror),
+        "get_terminal_size" => ctx.new_function(os_get_terminal_size),
 
         "O_RDONLY" => ctx.new_int(libc::O_RDONLY),
         "O_WRONLY" => ctx.new_int(libc::O_WRONLY),
diff --git a/vm/src/stdlib/winapi.rs b/vm/src/stdlib/winapi.rs
index 11956445a..4ba177654 100644
--- a/vm/src/stdlib/winapi.rs
+++ b/vm/src/stdlib/winapi.rs
@@ -130,9 +130,8 @@ fn _winapi_CreateProcess(
     args: CreateProcessArgs,
     vm: &VirtualMachine,
 ) -> PyResult<(usize, usize, u32, u32)> {
-    use winbase::STARTUPINFOEXW;
-    let mut si: STARTUPINFOEXW = unsafe { std::mem::zeroed() };
-    si.StartupInfo.cb = std::mem::size_of::<STARTUPINFOEXW>() as _;
+    let mut si = winbase::STARTUPINFOEXW::default();
+    si.StartupInfo.cb = std::mem::size_of_val(&si) as _;
 
     macro_rules! si_attr {
         ($attr:ident, $t:ty) => {{
@@ -207,7 +206,7 @@ fn _winapi_CreateProcess(
                 | winbase::CREATE_UNICODE_ENVIRONMENT,
             env as _,
             current_dir,
-            &mut si as *mut STARTUPINFOEXW as _,
+            &mut si as *mut winbase::STARTUPINFOEXW as _,
             &mut procinfo,
         )
     };