Emit Newline and handle indentation for Magic commands

Author: dhruvmanila

parser/src/lexer.rs

@@ -655,6 +655,12 @@ where
             // Detect indentation levels
             if self.at_begin_of_line {
                 self.handle_indentations()?;
+                if self.mode == Mode::Jupyter
+                    // https://github.com/ipython/ipython/blob/635815e8f1ded5b764d66cacc80bbe25e9e2587f/IPython/core/inputtransformer2.py#L345
+                    && matches!(self.window[0], Some('%' | '!' | '?' | '/' | ';' | ','))
+                {
+                    self.lex_and_emit_magic_command();
+                }
             }
 
             self.consume_normal()?;
@@ -703,10 +709,6 @@ where
                     spaces = 0;
                     tabs = 0;
                 }
-                // https://github.com/ipython/ipython/blob/635815e8f1ded5b764d66cacc80bbe25e9e2587f/IPython/core/inputtransformer2.py#L345
-                Some('%' | '!' | '?' | '/' | ';' | ',') if self.mode == Mode::Jupyter => {
-                    self.lex_and_emit_magic_command();
-                }
                 Some('\x0C') => {
                     // Form feed character!
                     // Reset indentation for the Emacs user.
@@ -1485,10 +1487,13 @@ mod tests {
         let tokens = lex_jupyter_source(&source);
         assert_eq!(
             tokens,
-            vec![Tok::MagicCommand {
-                value: "matplotlib   --inline".to_string(),
-                kind: MagicKind::Magic
-            },]
+            vec![
+                Tok::MagicCommand {
+                    value: "matplotlib   --inline".to_string(),
+                    kind: MagicKind::Magic
+                },
+                Tok::Newline
+            ]
         )
     }
 
@@ -1512,10 +1517,13 @@ mod tests {
         let tokens = lex_jupyter_source(&source);
         assert_eq!(
             tokens,
-            vec![Tok::MagicCommand {
-                value: "matplotlib ".to_string(),
-                kind: MagicKind::Magic
-            },]
+            vec![
+                Tok::MagicCommand {
+                    value: "matplotlib ".to_string(),
+                    kind: MagicKind::Magic
+                },
+                Tok::Newline
+            ]
         )
     }
 
@@ -1545,54 +1553,47 @@ mod tests {
                     value: "".to_string(),
                     kind: MagicKind::Magic,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Magic2,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Shell,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::ShCap,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Help,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Help2,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Paren,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Quote,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "".to_string(),
                     kind: MagicKind::Quote2,
                 },
+                Tok::Newline,
             ]
         )
     }
@@ -1611,10 +1612,8 @@ mod tests {
 !!cd /Users/foo/Library/Application\ Support/
 /foo 1 2
 ,foo 1 2
-;foo 1 2
-    !ls
-"
-        .trim();
+;foo 1 2"
+            .trim();
         let tokens = lex_jupyter_source(source);
         assert_eq!(
             tokens,
@@ -1623,66 +1622,78 @@ mod tests {
                     value: "foo".to_string(),
                     kind: MagicKind::Help,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "foo".to_string(),
                     kind: MagicKind::Help2,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "timeit a = b".to_string(),
                     kind: MagicKind::Magic,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "timeit a % 3".to_string(),
                     kind: MagicKind::Magic,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "matplotlib     --inline".to_string(),
                     kind: MagicKind::Magic,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "pwd   && ls -a | sed 's/^/\\\\    /'".to_string(),
                     kind: MagicKind::Shell,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "cd /Users/foo/Library/Application\\ Support/".to_string(),
                     kind: MagicKind::ShCap,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "foo 1 2".to_string(),
                     kind: MagicKind::Paren,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "foo 1 2".to_string(),
                     kind: MagicKind::Quote,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
                 Tok::MagicCommand {
                     value: "foo 1 2".to_string(),
                     kind: MagicKind::Quote2,
                 },
-                #[cfg(feature = "full-lexer")]
-                Tok::NonLogicalNewline,
+                Tok::Newline,
+            ]
+        )
+    }
+
+    #[test]
+    fn test_jupyter_magic_indentation() {
+        let source = r"
+if True:
+    %matplotlib \
+        --inline"
+            .trim();
+        let tokens = lex_jupyter_source(source);
+        assert_eq!(
+            tokens,
+            vec![
+                Tok::If,
+                Tok::True,
+                Tok::Colon,
+                Tok::Newline,
+                Tok::Indent,
                 Tok::MagicCommand {
-                    value: "ls".to_string(),
-                    kind: MagicKind::Shell,
+                    value: "matplotlib         --inline".to_string(),
+                    kind: MagicKind::Magic,
                 },
+                Tok::Newline,
+                Tok::Dedent,
             ]
         )
     }