Add line magic stmt and expr AST nodes (#31)

This PR introduces two new nodes for the parser to recognize the
`MagicCommand` tokens:

* `StmtLineMagic` for statement position i.e., magic commands on their
own line:
	```python
	%matplotlib inline
	!pwd
	```

* `ExprLineMagic` for expression position i.e., magic commands in an
assignment statement:
	```python
	# Only `?` and `!` are valid in this position
	dir = !pwd
	```

Both nodes are identical in their structure as in it contains the magic
kind and the command value as `String`.

Author: dhruvmanila

ast/src/generic.rs

@@ -1,7 +1,7 @@
 #![allow(clippy::derive_partial_eq_without_eq)]
 use crate::text_size::TextRange;
 pub use crate::{builtin::*, text_size::TextSize, ConversionFlag, Node};
-use std::fmt::Debug;
+use std::fmt::{self, Debug};
 
 // This file was originally generated from asdl by a python script, but we now edit it manually
 
@@ -317,6 +317,32 @@ pub enum Stmt {
     Break(StmtBreak),
     #[is(name = "continue_stmt")]
     Continue(StmtContinue),
+
+    // Jupyter notebook specific
+    #[is(name = "line_magic_stmt")]
+    LineMagic(StmtLineMagic),
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct StmtLineMagic {
+    pub range: TextRange,
+    pub kind: MagicKind,
+    pub value: String,
+}
+
+impl Node for StmtLineMagic {
+    const NAME: &'static str = "LineMagic";
+    const FIELD_NAMES: &'static [&'static str] = &["kind", "value"];
+}
+impl From<StmtLineMagic> for Stmt {
+    fn from(payload: StmtLineMagic) -> Self {
+        Stmt::LineMagic(payload)
+    }
+}
+impl From<StmtLineMagic> for Ast {
+    fn from(payload: StmtLineMagic) -> Self {
+        Stmt::from(payload).into()
+    }
 }
 
 /// See also [FunctionDef](https://docs.python.org/3/library/ast.html#ast.FunctionDef)
@@ -1084,6 +1110,32 @@ pub enum Expr {
     Tuple(ExprTuple),
     #[is(name = "slice_expr")]
     Slice(ExprSlice),
+
+    // Jupyter notebook specific
+    #[is(name = "line_magic_expr")]
+    LineMagic(ExprLineMagic),
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct ExprLineMagic {
+    pub range: TextRange,
+    pub kind: MagicKind,
+    pub value: String,
+}
+
+impl Node for ExprLineMagic {
+    const NAME: &'static str = "LineMagic";
+    const FIELD_NAMES: &'static [&'static str] = &["kind", "value"];
+}
+impl From<ExprLineMagic> for Expr {
+    fn from(payload: ExprLineMagic) -> Self {
+        Expr::LineMagic(payload)
+    }
+}
+impl From<ExprLineMagic> for Ast {
+    fn from(payload: ExprLineMagic) -> Self {
+        Expr::from(payload).into()
+    }
 }
 
 /// See also [BoolOp](https://docs.python.org/3/library/ast.html#ast.BoolOp)
@@ -3358,3 +3410,90 @@ impl Arguments {
         (args, with_defaults)
     }
 }
+
+/// The kind of magic command as defined in [IPython Syntax] in the IPython codebase.
+///
+/// [IPython Syntax]: https://github.com/ipython/ipython/blob/635815e8f1ded5b764d66cacc80bbe25e9e2587f/IPython/core/inputtransformer2.py#L335-L343
+#[derive(PartialEq, Eq, Debug, Clone, Hash, Copy)]
+pub enum MagicKind {
+    /// Send line to underlying system shell.
+    Shell,
+    /// Send line to system shell and capture output.
+    ShCap,
+    /// Show help on object.
+    Help,
+    /// Show help on object, with extra verbosity.
+    Help2,
+    /// Call magic function.
+    Magic,
+    /// Call cell magic function.
+    Magic2,
+    /// Call first argument with rest of line as arguments after splitting on whitespace
+    /// and quote each as string.
+    Quote,
+    /// Call first argument with rest of line as an argument quoted as a single string.
+    Quote2,
+    /// Call first argument with rest of line as arguments.
+    Paren,
+}
+
+impl TryFrom<char> for MagicKind {
+    type Error = String;
+
+    fn try_from(ch: char) -> Result<Self, Self::Error> {
+        match ch {
+            '!' => Ok(MagicKind::Shell),
+            '?' => Ok(MagicKind::Help),
+            '%' => Ok(MagicKind::Magic),
+            ',' => Ok(MagicKind::Quote),
+            ';' => Ok(MagicKind::Quote2),
+            '/' => Ok(MagicKind::Paren),
+            _ => Err(format!("Unexpected magic escape: {ch}")),
+        }
+    }
+}
+
+impl TryFrom<[char; 2]> for MagicKind {
+    type Error = String;
+
+    fn try_from(ch: [char; 2]) -> Result<Self, Self::Error> {
+        match ch {
+            ['!', '!'] => Ok(MagicKind::ShCap),
+            ['?', '?'] => Ok(MagicKind::Help2),
+            ['%', '%'] => Ok(MagicKind::Magic2),
+            [c1, c2] => Err(format!("Unexpected magic escape: {c1}{c2}")),
+        }
+    }
+}
+
+impl fmt::Display for MagicKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            MagicKind::Shell => f.write_str("!"),
+            MagicKind::ShCap => f.write_str("!!"),
+            MagicKind::Help => f.write_str("?"),
+            MagicKind::Help2 => f.write_str("??"),
+            MagicKind::Magic => f.write_str("%"),
+            MagicKind::Magic2 => f.write_str("%%"),
+            MagicKind::Quote => f.write_str(","),
+            MagicKind::Quote2 => f.write_str(";"),
+            MagicKind::Paren => f.write_str("/"),
+        }
+    }
+}
+
+impl MagicKind {
+    /// Returns the length of the magic command prefix.
+    pub fn prefix_len(self) -> TextSize {
+        let len = match self {
+            MagicKind::Shell
+            | MagicKind::Magic
+            | MagicKind::Help
+            | MagicKind::Quote
+            | MagicKind::Quote2
+            | MagicKind::Paren => 1,
+            MagicKind::ShCap | MagicKind::Magic2 | MagicKind::Help2 => 2,
+        };
+        len.into()
+    }
+}

ast/src/impls.rs

@@ -49,6 +49,7 @@ impl Expr {
             Expr::Lambda { .. } => "lambda",
             Expr::IfExp { .. } => "conditional expression",
             Expr::NamedExpr { .. } => "named expression",
+            Expr::LineMagic(_) => "line magic",
         }
     }
 }

ast/src/ranged.rs

@@ -207,6 +207,11 @@ impl Ranged for crate::generic::StmtContinue {
         self.range
     }
 }
+impl Ranged for crate::generic::StmtLineMagic {
+    fn range(&self) -> TextRange {
+        self.range
+    }
+}
 impl Ranged for crate::Stmt {
     fn range(&self) -> TextRange {
         match self {
@@ -238,6 +243,7 @@ impl Ranged for crate::Stmt {
             Self::Pass(node) => node.range(),
             Self::Break(node) => node.range(),
             Self::Continue(node) => node.range(),
+            Self::LineMagic(node) => node.range(),
         }
     }
 }
@@ -377,6 +383,11 @@ impl Ranged for crate::generic::ExprSlice {
         self.range
     }
 }
+impl Ranged for crate::generic::ExprLineMagic {
+    fn range(&self) -> TextRange {
+        self.range
+    }
+}
 impl Ranged for crate::Expr {
     fn range(&self) -> TextRange {
         match self {
@@ -407,6 +418,7 @@ impl Ranged for crate::Expr {
             Self::List(node) => node.range(),
             Self::Tuple(node) => node.range(),
             Self::Slice(node) => node.range(),
+            Self::LineMagic(node) => node.range(),
         }
     }
 }

core/src/mode.rs

@@ -15,16 +15,12 @@ pub enum Mode {
     ///
     /// ## Limitations:
     ///
-    /// These escaped commands are only supported when they are the only
-    /// statement on a line. If they're part of a larger statement such as
-    /// on the right-hand side of an assignment, the lexer will not recognize
-    /// them as escape commands.
-    ///
     /// For [Dynamic object information], the escape characters (`?`, `??`)
     /// must be used before an object. For example, `?foo` will be recognized,
     /// but `foo?` will not.
     ///
     /// ## Supported escape commands:
+    ///
     /// - [Magic command system] which is limited to [line magics] and can start
     ///   with `?` or `??`.
     /// - [Dynamic object information] which can start with `?` or `??`.