Make the REPL handle continuation better

Author: coolreader18

src/main.rs

@@ -159,6 +159,14 @@ fn get_history_path() -> PathBuf {
     xdg_dirs.place_cache_file("repl_history.txt").unwrap()
 }
 
+fn get_prompt(vm: &mut VirtualMachine, prompt_name: &str) -> String {
+    vm.sys_module
+        .get_attr(prompt_name)
+        .as_ref()
+        .map(objstr::get_value)
+        .unwrap_or_else(String::new)
+}
+
 fn run_shell(vm: &mut VirtualMachine) -> PyResult {
     println!(
         "Welcome to the magnificent Rust Python {} interpreter",
@@ -176,33 +184,45 @@ fn run_shell(vm: &mut VirtualMachine) -> PyResult {
         println!("No previous history.");
     }
 
-    let ps1 = &objstr::get_value(&vm.sys_module.get_attr("ps1").unwrap());
-    let ps2 = &objstr::get_value(&vm.sys_module.get_attr("ps2").unwrap());
-    let mut prompt = ps1;
+    let mut prompt = get_prompt(vm, "ps1");
+
+    let mut continuing = false;
 
     loop {
-        match repl.readline(prompt) {
+        if !continuing {
+            prompt = get_prompt(vm, "ps1");
+        }
+        match repl.readline(&prompt) {
             Ok(line) => {
                 debug!("You entered {:?}", line);
                 input.push_str(&line);
-                input.push_str("\n");
+                input.push('\n');
                 repl.add_history_entry(line.trim_end());
 
+                if continuing {
+                    if line.is_empty() {
+                        continuing = false;
+                    } else {
+                        continue;
+                    }
+                }
+
                 match shell_exec(vm, &input, vars.clone()) {
                     Err(CompileError::Parse(ParseError::EOF(_))) => {
-                        prompt = ps2;
+                        prompt = get_prompt(vm, "ps2");
+                        continuing = true;
                         continue;
                     }
                     _ => {
-                        prompt = ps1;
                         input = String::new();
                     }
                 }
             }
             Err(ReadlineError::Interrupted) => {
                 // TODO: Raise a real KeyboardInterrupt exception
                 println!("^C");
-                break;
+                continuing = false;
+                continue;
             }
             Err(ReadlineError::Eof) => {
                 break;