support format specs in f-strings

Author: OddCoincidence

parser/src/ast.rs

@@ -315,7 +315,14 @@ pub enum Number {
 
 #[derive(Debug, PartialEq)]
 pub enum StringGroup {
-    Constant { value: String },
-    FormattedValue { value: Box<Expression> },
-    Joined { values: Vec<StringGroup> },
+    Constant {
+        value: String,
+    },
+    FormattedValue {
+        value: Box<Expression>,
+        spec: String,
+    },
+    Joined {
+        values: Vec<StringGroup>,
+    },
 }

parser/src/parser.rs

@@ -86,8 +86,14 @@ impl From<FStringError>
 }
 
 enum ParseState {
-    Text { content: String },
-    FormattedValue { expression: String, depth: usize },
+    Text {
+        content: String,
+    },
+    FormattedValue {
+        expression: String,
+        spec: Option<String>,
+        depth: usize,
+    },
 }
 
 pub fn parse_fstring(source: &str) -> Result<ast::StringGroup, FStringError> {
@@ -114,6 +120,7 @@ pub fn parse_fstring(source: &str) -> Result<ast::StringGroup, FStringError> {
 
                         FormattedValue {
                             expression: String::new(),
+                            spec: None,
                             depth: 0,
                         }
                     }
@@ -135,17 +142,28 @@ pub fn parse_fstring(source: &str) -> Result<ast::StringGroup, FStringError> {
 
             FormattedValue {
                 mut expression,
+                mut spec,
                 depth,
             } => match ch {
+                ':' if depth == 0 => FormattedValue {
+                    expression,
+                    spec: Some(String::new()),
+                    depth,
+                },
                 '{' => {
                     if let Some('{') = chars.peek() {
                         expression.push_str("{{");
                         chars.next();
-                        FormattedValue { expression, depth }
+                        FormattedValue {
+                            expression,
+                            spec,
+                            depth,
+                        }
                     } else {
                         expression.push('{');
                         FormattedValue {
                             expression,
+                            spec,
                             depth: depth + 1,
                         }
                     }
@@ -154,28 +172,42 @@ pub fn parse_fstring(source: &str) -> Result<ast::StringGroup, FStringError> {
                     if let Some('}') = chars.peek() {
                         expression.push_str("}}");
                         chars.next();
-                        FormattedValue { expression, depth }
+                        FormattedValue {
+                            expression,
+                            spec,
+                            depth,
+                        }
                     } else if depth > 0 {
                         expression.push('}');
                         FormattedValue {
                             expression,
+                            spec,
                             depth: depth - 1,
                         }
                     } else {
                         values.push(ast::StringGroup::FormattedValue {
-                            value: Box::new(match parse_expression(dbg!(expression.trim())) {
+                            value: Box::new(match parse_expression(expression.trim()) {
                                 Ok(expr) => expr,
                                 Err(_) => return Err(FStringError::InvalidExpression),
                             }),
+                            spec: spec.unwrap_or_default(),
                         });
                         Text {
                             content: String::new(),
                         }
                     }
                 }
                 _ => {
-                    expression.push(ch);
-                    FormattedValue { expression, depth }
+                    if let Some(spec) = spec.as_mut() {
+                        spec.push(ch)
+                    } else {
+                        expression.push(ch);
+                    }
+                    FormattedValue {
+                        expression,
+                        spec,
+                        depth,
+                    }
                 }
             },
         };
@@ -617,10 +649,12 @@ mod tests {
             ast::StringGroup::Joined {
                 values: vec![
                     ast::StringGroup::FormattedValue {
-                        value: Box::new(mk_ident("a"))
+                        value: Box::new(mk_ident("a")),
+                        spec: String::new(),
                     },
                     ast::StringGroup::FormattedValue {
-                        value: Box::new(mk_ident("b"))
+                        value: Box::new(mk_ident("b")),
+                        spec: String::new(),
                     },
                     ast::StringGroup::Constant {
                         value: "{foo}".to_owned()

tests/snippets/fstrings.py

@@ -11,6 +11,7 @@
 assert f'{f"{{"}' == '{'
 assert f'{f"}}"}' == '}'
 assert f'{foo}' f"{foo}" 'foo' == 'barbarfoo'
-assert f'{"!:"}' == '!:'
+#assert f'{"!:"}' == '!:'
 #assert f"{1 != 2}" == 'True'
 assert fr'x={4*10}\n' == 'x=40\\n'
+assert f'{16:0>+#10x}' == '00000+0x10'

vm/src/bytecode.rs

@@ -167,7 +167,9 @@ pub enum Instruction {
         after: usize,
     },
     Unpack,
-    FormatValue,
+    FormatValue {
+        spec: String,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq)]

vm/src/compile.rs

@@ -1336,9 +1336,9 @@ impl Compiler {
                     },
                 });
             }
-            ast::StringGroup::FormattedValue { value } => {
+            ast::StringGroup::FormattedValue { value, spec } => {
                 self.compile_expression(value)?;
-                self.emit(Instruction::FormatValue);
+                self.emit(Instruction::FormatValue { spec: spec.clone() });
             }
         }
         Ok(())

vm/src/frame.rs

@@ -640,10 +640,11 @@ impl Frame {
                 }
                 Ok(None)
             }
-            bytecode::Instruction::FormatValue => {
+            bytecode::Instruction::FormatValue { spec } => {
                 let value = self.pop_value();
-                let formatted = vm.to_pystr(&value)?;
-                self.push_value(vm.new_str(formatted));
+                let spec = vm.new_str(spec.clone());
+                let formatted = vm.call_method(&value, "__format__", vec![spec])?;
+                self.push_value(formatted);
                 Ok(None)
             }
         }

vm/src/stdlib/ast.rs

@@ -570,7 +570,7 @@ fn string_to_ast(ctx: &PyContext, string: &ast::StringGroup) -> PyObjectRef {
             ctx.set_attr(&node, "s", ctx.new_str(value.clone()));
             node
         }
-        ast::StringGroup::FormattedValue { value } => {
+        ast::StringGroup::FormattedValue { value, .. } => {
             let node = create_node(ctx, "FormattedValue");
             let py_value = expression_to_ast(ctx, value);
             ctx.set_attr(&node, "value", py_value);