a few more tweaks

5 files changed, 56 insertions, 53 deletions

diff --git a/src/async_stdin.rs b/src/async_stdin.rs
index cf8040f..5cfea44 100644
--- a/src/async_stdin.rs
+++ b/src/async_stdin.rs
@@ -70,14 +70,13 @@ impl tokio::io::AsyncRead for Stdin {
         match self.input.poll_read_ready(ready)? {
             futures::Async::Ready(_) => {
                 let res = self.input.poll_read(buf);
+
                 // XXX i'm pretty sure this is wrong (if the single poll_read
                 // call didn't return all waiting data, clearing read ready
                 // state means that we won't get the rest until some more data
                 // beyond that appears), but i don't know that there's a way
                 // to do it correctly given that poll_read blocks
-                if let Err(e) = self.input.clear_read_ready(ready) {
-                    return Err(e);
-                }
+                self.input.clear_read_ready(ready)?;
 
                 res
             }
diff --git a/src/cmd/record.rs b/src/cmd/record.rs
index 4577687..25c11f7 100644
--- a/src/cmd/record.rs
+++ b/src/cmd/record.rs
@@ -158,7 +158,14 @@ impl RecordSession {
         match self.process.poll()? {
             futures::Async::Ready(Some(e)) => {
                 match e {
-                    crate::process::Event::CommandStart(..) => {}
+                    crate::process::Event::CommandStart(..) => {
+                        if self.raw_screen.is_none() {
+                            self.raw_screen = Some(
+                                crossterm::RawScreen::into_raw_mode()
+                                    .context(crate::error::ToRawMode)?,
+                            );
+                        }
+                    }
                     crate::process::Event::CommandExit(..) => {
                         self.done = true;
                     }
@@ -262,12 +269,6 @@ impl futures::future::Future for RecordSession {
     type Error = Error;
 
     fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
-        if self.raw_screen.is_none() {
-            self.raw_screen = Some(
-                crossterm::RawScreen::into_raw_mode()
-                    .context(crate::error::ToRawMode)?,
-            );
-        }
         crate::component_future::poll_future(self, Self::POLL_FNS)
     }
 }
diff --git a/src/cmd/stream.rs b/src/cmd/stream.rs
index dac3d2c..1e60a25 100644
--- a/src/cmd/stream.rs
+++ b/src/cmd/stream.rs
@@ -232,7 +232,14 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
         match self.process.poll()? {
             futures::Async::Ready(Some(e)) => {
                 match e {
-                    crate::process::Event::CommandStart(..) => {}
+                    crate::process::Event::CommandStart(..) => {
+                        if self.raw_screen.is_none() {
+                            self.raw_screen = Some(
+                                crossterm::RawScreen::into_raw_mode()
+                                    .context(crate::error::ToRawMode)?,
+                            );
+                        }
+                    }
                     crate::process::Event::CommandExit(..) => {
                         self.done = true;
                     }
@@ -330,12 +337,6 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
     type Error = Error;
 
     fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
-        if self.raw_screen.is_none() {
-            self.raw_screen = Some(
-                crossterm::RawScreen::into_raw_mode()
-                    .context(crate::error::ToRawMode)?,
-            );
-        }
         crate::component_future::poll_future(self, Self::POLL_FNS)
     }
 }
diff --git a/src/cmd/watch.rs b/src/cmd/watch.rs
index 3129bba..6c69b14 100644
--- a/src/cmd/watch.rs
+++ b/src/cmd/watch.rs
@@ -527,6 +527,18 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
     ];
 
     fn poll_input(&mut self) -> Result<crate::component_future::Poll<()>> {
+        if self.raw_screen.is_none() {
+            self.raw_screen = Some(
+                crossterm::RawScreen::into_raw_mode()
+                    .context(crate::error::ToRawMode)?,
+            );
+        }
+        if let State::Temporary = self.state {
+            self.state = State::LoggingIn {
+                alternate_screen: new_alternate_screen()?,
+            }
+        }
+
         match self.key_reader.poll()? {
             futures::Async::Ready(Some(e)) => {
                 let quit = match &mut self.state {
@@ -626,17 +638,6 @@ impl<S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static>
     type Error = Error;
 
     fn poll(&mut self) -> futures::Poll<Self::Item, Self::Error> {
-        if let State::Temporary = self.state {
-            self.state = State::LoggingIn {
-                alternate_screen: new_alternate_screen()?,
-            }
-        }
-        if self.raw_screen.is_none() {
-            self.raw_screen = Some(
-                crossterm::RawScreen::into_raw_mode()
-                    .context(crate::error::ToRawMode)?,
-            );
-        }
         let res = crate::component_future::poll_future(self, Self::POLL_FNS);
         if res.is_err() {
             self.state = State::Temporary; // drop alternate screen
diff --git a/src/process.rs b/src/process.rs
index 3756837..70ad256 100644
--- a/src/process.rs
+++ b/src/process.rs
@@ -117,6 +117,32 @@ impl<R: tokio::io::AsyncRead + 'static> Process<R> {
             return Ok(crate::component_future::Poll::NothingToDo);
         }
 
+        if self.state.pty.is_none() {
+            self.state.pty = Some(
+                tokio_pty_process::AsyncPtyMaster::open()
+                    .context(crate::error::OpenPty)?,
+            );
+            log::debug!(
+                "openpty({})",
+                self.state.pty.as_ref().unwrap().as_raw_fd()
+            );
+        }
+
+        if self.state.process.is_none() {
+            self.state.process = Some(
+                std::process::Command::new(&self.cmd)
+                    .args(&self.args)
+                    .spawn_pty_async(self.state.pty())
+                    .context(crate::error::SpawnProcess {
+                        cmd: self.cmd.clone(),
+                    })?,
+            );
+            log::debug!(
+                "spawn({})",
+                self.state.process.as_ref().unwrap().id()
+            );
+        }
+
         self.started = true;
         Ok(crate::component_future::Poll::Event(Event::CommandStart(
             self.cmd.clone(),
@@ -250,31 +276,6 @@ impl<R: tokio::io::AsyncRead + 'static> futures::stream::Stream
     type Error = Error;
 
     fn poll(&mut self) -> futures::Poll<Option<Self::Item>, Self::Error> {
-        if self.state.pty.is_none() {
-            self.state.pty = Some(
-                tokio_pty_process::AsyncPtyMaster::open()
-                    .context(crate::error::OpenPty)?,
-            );
-            log::debug!(
-                "openpty({})",
-                self.state.pty.as_ref().unwrap().as_raw_fd()
-            );
-        }
-        if self.state.process.is_none() {
-            self.state.process = Some(
-                std::process::Command::new(&self.cmd)
-                    .args(&self.args)
-                    .spawn_pty_async(self.state.pty())
-                    .context(crate::error::SpawnProcess {
-                        cmd: self.cmd.clone(),
-                    })?,
-            );
-            log::debug!(
-                "spawn({})",
-                self.state.process.as_ref().unwrap().id()
-            );
-        }
-
         crate::component_future::poll_stream(self, Self::POLL_FNS)
     }
 }