Implement `attach --create` subcommand flag to create session if one does not exist (#731)

* Implement attach --create

* fixup! Implement attach --create

* fixup! Implement attach --create

* fixup! Implement attach --create

* fixup! Implement attach --create

3 files changed, 100 insertions, 24 deletions

diff --git a/src/main.rs b/src/main.rs
index 294464504..0a91ca652 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,10 @@ mod sessions;
 mod tests;
 
 use crate::install::populate_data_dir;
-use sessions::{assert_session, assert_session_ne, get_active_session, list_sessions};
+use sessions::{
+    assert_session, assert_session_ne, get_active_session, get_sessions, list_sessions,
+    print_sessions, session_exists, ActiveSession,
+};
 use std::process;
 use zellij_client::{os_input_output::get_client_os_input, start_client, ClientInfo};
 use zellij_server::{os_input_output::get_server_os_input, start_server};
@@ -52,29 +55,75 @@ pub fn main() {
             }
         };
         if let Some(Command::Sessions(Sessions::Attach {
-            mut session_name,
+            session_name,
             force,
+            create,
             options,
         })) = opts.command.clone()
         {
-            if let Some(session) = session_name.as_ref() {
-                assert_session(session);
-            } else {
-                session_name = Some(get_active_session());
-            }
-
             let config_options = match options {
                 Some(SessionCommand::Options(o)) => config_options.merge(o),
                 None => config_options,
             };
 
+            let (client, attach_layout) = match session_name.as_ref() {
+                Some(session) => {
+                    if create {
+                        if !session_exists(session).unwrap() {
+                            (ClientInfo::New(session_name.unwrap()), layout)
+                        } else {
+                            (
+                                ClientInfo::Attach(
+                                    session_name.unwrap(),
+                                    force,
+                                    config_options.clone(),
+                                ),
+                                None,
+                            )
+                        }
+                    } else {
+                        assert_session(session);
+                        (
+                            ClientInfo::Attach(
+                                session_name.unwrap(),
+                                force,
+                                config_options.clone(),
+                            ),
+                            None,
+                        )
+                    }
+                }
+                None => match get_active_session() {
+                    ActiveSession::None => {
+                        if create {
+                            (
+                                ClientInfo::New(names::Generator::default().next().unwrap()),
+                                layout,
+                            )
+                        } else {
+                            println!("No active zellij sessions found.");
+                            process::exit(1);
+                        }
+                    }
+                    ActiveSession::One(session_name) => (
+                        ClientInfo::Attach(session_name, force, config_options.clone()),
+                        None,
+                    ),
+                    ActiveSession::Many => {
+                        println!("Please specify the session name to attach to. The following sessions are active:");
+                        print_sessions(get_sessions().unwrap());
+                        process::exit(1);
+                    }
+                },
+            };
+
             start_client(
                 Box::new(os_input),
                 opts,
                 config,
-                config_options.clone(),
-                ClientInfo::Attach(session_name.unwrap(), force, config_options),
-                None,
+                config_options,
+                client,
+                attach_layout,
             );
         } else {
             let session_name = opts
diff --git a/src/sessions.rs b/src/sessions.rs
index fb1a96352..9f1e20f0a 100644
--- a/src/sessions.rs
+++ b/src/sessions.rs
@@ -6,7 +6,7 @@ use zellij_utils::{
     ipc::{ClientToServerMsg, IpcSenderWithContext},
 };
 
-fn get_sessions() -> Result<Vec<String>, io::ErrorKind> {
+pub(crate) fn get_sessions() -> Result<Vec<String>, io::ErrorKind> {
     match fs::read_dir(&*ZELLIJ_SOCK_DIR) {
         Ok(files) => {
             let mut sessions = Vec::new();
@@ -47,7 +47,7 @@ fn assert_socket(name: &str) -> bool {
     }
 }
 
-fn print_sessions(sessions: Vec<String>) {
+pub(crate) 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 {
@@ -59,22 +59,29 @@ fn print_sessions(sessions: Vec<String>) {
     })
 }
 
-pub(crate) fn get_active_session() -> String {
+pub(crate) enum ActiveSession {
+    None,
+    One(String),
+    Many,
+}
+
+pub(crate) fn get_active_session() -> ActiveSession {
     match get_sessions() {
         Ok(mut sessions) => {
             if sessions.len() == 1 {
-                return sessions.pop().unwrap();
+                return ActiveSession::One(sessions.pop().unwrap());
             }
             if sessions.is_empty() {
-                println!("No active zellij sessions found.");
+                ActiveSession::None
             } else {
-                println!("Please specify the session name to attach to. The following sessions are active:");
-                print_sessions(sessions);
+                ActiveSession::Many
             }
         }
-        Err(e) => eprintln!("Error occured: {:?}", e),
+        Err(e) => {
+            eprintln!("Error occured: {:?}", e);
+            process::exit(1);
+        }
     }
-    process::exit(1);
 }
 
 pub(crate) fn list_sessions() {
@@ -95,15 +102,30 @@ pub(crate) fn list_sessions() {
     process::exit(exit_code);
 }
 
-pub(crate) fn assert_session(name: &str) {
-    match get_sessions() {
+pub(crate) fn session_exists(name: &str) -> Result<bool, io::ErrorKind> {
+    return match get_sessions() {
         Ok(sessions) => {
             if sessions.iter().any(|s| s == name) {
+                return Ok(true);
+            }
+            Ok(false)
+        }
+        Err(e) => Err(e),
+    };
+}
+
+pub(crate) fn assert_session(name: &str) {
+    match session_exists(name) {
+        Ok(result) => {
+            if result {
                 return;
+            } else {
+                println!("No session named {:?} found.", name);
             }
-            println!("No session named {:?} found.", name);
         }
-        Err(e) => eprintln!("Error occured: {:?}", e),
+        Err(e) => {
+            eprintln!("Error occured: {:?}", e);
+        }
     };
     process::exit(1);
 }
diff --git a/zellij-utils/src/cli.rs b/zellij-utils/src/cli.rs
index a8dc9c0aa..eb52589ae 100644
--- a/zellij-utils/src/cli.rs
+++ b/zellij-utils/src/cli.rs
@@ -85,6 +85,11 @@ pub enum Sessions {
         /// zellij client (if any) and attach to this.
         #[structopt(long, short)]
         force: bool,
+
+        /// Create a session if one does not exist.
+        #[structopt(short, long)]
+        create: bool,
+
         /// Change the behaviour of zellij
         #[structopt(subcommand, name = "options")]
         options: Option<SessionCommand>,