From 0bd05cbcb46d607079608bb83bdefba97bf7a764 Mon Sep 17 00:00:00 2001
From: Kunal Mohan <kunalmohan99@gmail.com>
Date: Sat, 29 May 2021 00:32:06 +0530
Subject: [PATCH] Make session-name option in attach command If only one
 session is running, attach to it. Otherwise list the active sessions (if any)

---
 src/main.rs              | 12 +++--
 src/sessions.rs          | 97 +++++++++++++++++++++++-----------------
 zellij-server/src/lib.rs |  3 +-
 zellij-utils/src/cli.rs  |  2 +-
 4 files changed, 65 insertions(+), 49 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index efb991567d..b9147fab2e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,7 @@ mod sessions;
 mod tests;
 
 use crate::install::populate_data_dir;
-use sessions::{assert_session, assert_session_ne, list_sessions};
+use sessions::{assert_session, assert_session_ne, get_active_session, list_sessions};
 use std::convert::TryFrom;
 use std::process;
 use zellij_client::{os_input_output::get_client_os_input, start_client, ClientInfo};
@@ -54,16 +54,20 @@ pub fn main() {
             }
         };
         if let Some(Command::Sessions(Sessions::Attach {
-            session_name,
+            mut session_name,
             force,
         })) = opts.command.clone()
         {
-            assert_session(&session_name);
+            if let Some(session) = session_name.as_ref() {
+                assert_session(session);
+            } else {
+                session_name = Some(get_active_session());
+            }
             start_client(
                 Box::new(os_input),
                 opts,
                 config,
-                ClientInfo::Attach(session_name, force),
+                ClientInfo::Attach(session_name.unwrap(), force),
             );
         } else {
             let session_name = opts
diff --git a/src/sessions.rs b/src/sessions.rs
index fd834e4080..fb1a963525 100644
--- a/src/sessions.rs
+++ b/src/sessions.rs
@@ -29,22 +29,61 @@ fn get_sessions() -> Result<Vec<String>, io::ErrorKind> {
     }
 }
 
+fn assert_socket(name: &str) -> bool {
+    let path = &*ZELLIJ_SOCK_DIR.join(name);
+    match LocalSocketStream::connect(path) {
+        Ok(stream) => {
+            IpcSenderWithContext::new(stream).send(ClientToServerMsg::ClientExited);
+            true
+        }
+        Err(e) => {
+            if e.kind() == io::ErrorKind::ConnectionRefused {
+                drop(fs::remove_file(path));
+                false
+            } else {
+                true
+            }
+        }
+    }
+}
+
+fn print_sessions(sessions: Vec<String>) {
+    let curr_session = std::env::var("ZELLIJ_SESSION_NAME").unwrap_or_else(|_| "".into());
+    sessions.iter().for_each(|session| {
+        let suffix = if curr_session == *session {
+            " (current)"
+        } else {
+            ""
+        };
+        println!("{}{}", session, suffix);
+    })
+}
+
+pub(crate) fn get_active_session() -> String {
+    match get_sessions() {
+        Ok(mut sessions) => {
+            if sessions.len() == 1 {
+                return sessions.pop().unwrap();
+            }
+            if sessions.is_empty() {
+                println!("No active zellij sessions found.");
+            } else {
+                println!("Please specify the session name to attach to. The following sessions are active:");
+                print_sessions(sessions);
+            }
+        }
+        Err(e) => eprintln!("Error occured: {:?}", e),
+    }
+    process::exit(1);
+}
+
 pub(crate) fn list_sessions() {
     let exit_code = match get_sessions() {
         Ok(sessions) => {
             if sessions.is_empty() {
                 println!("No active zellij sessions found.");
             } else {
-                let curr_session =
-                    std::env::var("ZELLIJ_SESSION_NAME").unwrap_or_else(|_| "".into());
-                sessions.iter().for_each(|session| {
-                    let suffix = if curr_session == *session {
-                        " (current)"
-                    } else {
-                        ""
-                    };
-                    println!("{}{}", session, suffix);
-                })
+                print_sessions(sessions);
             }
             0
         }
@@ -57,53 +96,27 @@ pub(crate) fn list_sessions() {
 }
 
 pub(crate) fn assert_session(name: &str) {
-    let exit_code = match get_sessions() {
+    match get_sessions() {
         Ok(sessions) => {
             if sessions.iter().any(|s| s == name) {
                 return;
             }
             println!("No session named {:?} found.", name);
-            0
-        }
-        Err(e) => {
-            eprintln!("Error occured: {:?}", e);
-            1
         }
+        Err(e) => eprintln!("Error occured: {:?}", e),
     };
-    process::exit(exit_code);
+    process::exit(1);
 }
 
 pub(crate) fn assert_session_ne(name: &str) {
-    let exit_code = match get_sessions() {
+    match get_sessions() {
         Ok(sessions) => {
             if sessions.iter().all(|s| s != name) {
                 return;
             }
             println!("Session with name {:?} aleady exists. Use attach command to connect to it or specify a different name.", name);
-            0
-        }
-        Err(e) => {
-            eprintln!("Error occured: {:?}", e);
-            1
         }
+        Err(e) => eprintln!("Error occured: {:?}", e),
     };
-    process::exit(exit_code);
-}
-
-fn assert_socket(name: &str) -> bool {
-    let path = &*ZELLIJ_SOCK_DIR.join(name);
-    match LocalSocketStream::connect(path) {
-        Ok(stream) => {
-            IpcSenderWithContext::new(stream).send(ClientToServerMsg::ClientExited);
-            true
-        }
-        Err(e) => {
-            if e.kind() == io::ErrorKind::ConnectionRefused {
-                drop(fs::remove_file(path));
-                false
-            } else {
-                true
-            }
-        }
-    }
+    process::exit(1);
 }
diff --git a/zellij-server/src/lib.rs b/zellij-server/src/lib.rs
index 2f0389ef7a..b40362001c 100644
--- a/zellij-server/src/lib.rs
+++ b/zellij-server/src/lib.rs
@@ -27,8 +27,7 @@ use crate::{
 };
 use route::route_thread_main;
 use zellij_utils::{
-    channels,
-    channels::{ChannelWithContext, SenderWithContext},
+    channels::{self, ChannelWithContext, SenderWithContext},
     cli::CliArgs,
     errors::{ContextType, ErrorInstruction, ServerContext},
     input::{get_mode_info, options::Options},
diff --git a/zellij-utils/src/cli.rs b/zellij-utils/src/cli.rs
index bf7851c9f5..2ac0e986ba 100644
--- a/zellij-utils/src/cli.rs
+++ b/zellij-utils/src/cli.rs
@@ -72,7 +72,7 @@ pub enum Sessions {
     #[structopt(alias = "a")]
     Attach {
         /// Name of the session to attach to.
-        session_name: String,
+        session_name: Option<String>,
 
         /// Force attach- session will detach from the other
         /// zellij client (if any) and attach to this.