Fix FORMAT_VALUE for values that have both a conversion and a format_spec.

zrax · zrax · commit 0b45b5fa076d · 2024-08-01T13:28:43.000-07:00
Also output the conversion and flags in disassembly.
diff --git a/ASTNode.h b/ASTNode.h
@@ -689,7 +689,9 @@ class ASTFormattedValue : public ASTNode {
         STR = 1,
         REPR = 2,
         ASCII = 3,
-        FMTSPEC = 4
+        CONVERSION_MASK = 0x03,
+
+        HAVE_FMT_SPEC = 4,
     };
 
     ASTFormattedValue(PycRef<ASTNode> val, ConversionFlag conversion,
diff --git a/ASTree.cpp b/ASTree.cpp
@@ -978,29 +978,14 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
         case Pyc::FORMAT_VALUE_A:
             {
                 auto conversion_flag = static_cast<ASTFormattedValue::ConversionFlag>(operand);
-                switch (conversion_flag) {
-                case ASTFormattedValue::ConversionFlag::NONE:
-                case ASTFormattedValue::ConversionFlag::STR:
-                case ASTFormattedValue::ConversionFlag::REPR:
-                case ASTFormattedValue::ConversionFlag::ASCII:
-                    {
-                        auto val = stack.top();
-                        stack.pop();
-                        stack.push(new ASTFormattedValue(val, conversion_flag, nullptr));
-                    }
-                    break;
-                case ASTFormattedValue::ConversionFlag::FMTSPEC:
-                    {
-                        auto format_spec = stack.top();
-                        stack.pop();
-                        auto val = stack.top();
-                        stack.pop();
-                        stack.push(new ASTFormattedValue(val, conversion_flag, format_spec));
-                    }
-                    break;
-                default:
-                    fprintf(stderr, "Unsupported FORMAT_VALUE_A conversion flag: %d\n", operand);
+                PycRef<ASTNode> format_spec = nullptr;
+                if (conversion_flag & ASTFormattedValue::HAVE_FMT_SPEC) {
+                    format_spec = stack.top();
+                    stack.pop();
                 }
+                auto val = stack.top();
+                stack.pop();
+                stack.push(new ASTFormattedValue(val, conversion_flag, format_spec));
             }
             break;
         case Pyc::GET_AWAITABLE:
@@ -2659,23 +2644,21 @@ void print_formatted_value(PycRef<ASTFormattedValue> formatted_value, PycModule*
     pyc_output << "{";
     print_src(formatted_value->val(), mod, pyc_output);
 
-    switch (formatted_value->conversion()) {
-    case ASTFormattedValue::ConversionFlag::NONE:
+    switch (formatted_value->conversion() & ASTFormattedValue::CONVERSION_MASK) {
+    case ASTFormattedValue::NONE:
         break;
-    case ASTFormattedValue::ConversionFlag::STR:
+    case ASTFormattedValue::STR:
         pyc_output << "!s";
         break;
-    case ASTFormattedValue::ConversionFlag::REPR:
+    case ASTFormattedValue::REPR:
         pyc_output << "!r";
         break;
-    case ASTFormattedValue::ConversionFlag::ASCII:
+    case ASTFormattedValue::ASCII:
         pyc_output << "!a";
         break;
-    case ASTFormattedValue::ConversionFlag::FMTSPEC:
+    }
+    if (formatted_value->conversion() & ASTFormattedValue::HAVE_FMT_SPEC) {
         pyc_output << ":" << formatted_value->format_spec().cast<ASTObject>()->object().cast<PycString>()->value();
-        break;
-    default:
-        fprintf(stderr, "Unsupported NODE_FORMATTEDVALUE conversion flag: %d\n", formatted_value->conversion());
     }
     pyc_output << "}";
 }
diff --git a/bytecode.cpp b/bytecode.cpp
@@ -331,6 +331,11 @@ void bc_disasm(std::ostream& pyc_output, PycRef<PycCode> code, PycModule* mod,
     };
     static const size_t intrinsic2_names_len = sizeof(intrinsic2_names) / sizeof(intrinsic2_names[0]);
 
+    static const char *format_value_names[] = {
+        "FVC_NONE", "FVC_STR", "FVC_REPR", "FVC_ASCII",
+    };
+    static const size_t format_value_names_len = sizeof(format_value_names) / sizeof(format_value_names[0]);
+
     PycBuffer source(code->code()->value(), code->code()->length());
 
     int opcode, operand;
@@ -530,6 +535,18 @@ void bc_disasm(std::ostream& pyc_output, PycRef<PycCode> code, PycModule* mod,
                 else
                     formatted_print(pyc_output, "%d (UNKNOWN)", operand);
                 break;
+            case Pyc::FORMAT_VALUE_A:
+                {
+                    auto conv = static_cast<size_t>(operand & 0x03);
+                    const char *flag = (operand & 0x04) ? " | FVS_HAVE_SPEC" : "";
+                    if (conv < format_value_names_len) {
+                        formatted_print(pyc_output, "%d (%s%s)", operand,
+                                        format_value_names[conv], flag);
+                    } else {
+                        formatted_print(pyc_output, "%d (UNKNOWN)", operand);
+                    }
+                }
+                break;
             default:
                 formatted_print(pyc_output, "%d", operand);
                 break;
diff --git a/bytecode_ops.inl b/bytecode_ops.inl
@@ -204,7 +204,7 @@ OPCODE_A(BUILD_MAP_UNPACK_WITH_CALL)    // Python 3.5                   A=(count
 OPCODE_A(BUILD_TUPLE_UNPACK)            // Python 3.5 - 3.8             A=count
 OPCODE_A(BUILD_SET_UNPACK)              // Python 3.5 - 3.8             A=count
 OPCODE_A(SETUP_ASYNC_WITH)              // Python 3.5 - 3.10            rel jmp +A
-OPCODE_A(FORMAT_VALUE)                  // Python 3.6 ->                A=conversion_type
+OPCODE_A(FORMAT_VALUE)                  // Python 3.6 ->                A=(conversion_type&0x3)+(flags)
 OPCODE_A(BUILD_CONST_KEY_MAP)           // Python 3.6 ->                A=count
 OPCODE_A(BUILD_STRING)                  // Python 3.6 ->                A=count
 OPCODE_A(BUILD_TUPLE_UNPACK_WITH_CALL)  // Python 3.6 - 3.8             A=count
diff --git a/tests/compiled/f-string.3.7.pyc b/tests/compiled/f-string.3.7.pyc
diff --git a/tests/input/f-string.py b/tests/input/f-string.py
@@ -28,10 +28,11 @@
 print(f'''some {{braces}} {"inner literal: {braces} {{double braces}}"}''')
 print(f'''f-string dict {some_dict[2]} and {{function call in expression}}: {max([1,20,3])}''')
 print(f'{(lambda x: x*2)(3)}')
+print(f'{var3!s:4.5}')
 msg = (
     f'a {var1}'
     f'cool'
     f'multiline {var2}\n'
     f'f-string {var3}'
 )
-print(f'{now:%Y-%m-%d %H:%M}')
+print(f'{now:%Y-%m-%d %H:%M}')
diff --git a/tests/tokenized/f-string.txt b/tests/tokenized/f-string.txt
@@ -27,5 +27,6 @@ print ( f'{var3 * x} {var3:.2f} {var3:.5f} {x:02} {x * x:3} {x * x * x:4} {s1:>1
 print ( f'some {{braces}} {\'inner literal: {braces} {{double braces}}\'}' ) <EOL>
 print ( f'f-string dict {some_dict[2]} and {{function call in expression}}: {max([\n    1,\n    20,\n    3])}' ) <EOL>
 print ( f'{(lambda x: x * 2)(3)}' ) <EOL>
+print ( f'{var3!s:4.5}' ) <EOL>
 msg = f'a {var1}coolmultiline {var2}\nf-string {var3}' <EOL>
 print ( f'{now:%Y-%m-%d %H:%M}' ) <EOL>
