Handle SET_FUNCTION_ATTRIBUTE correctly

It leaves the function on the stack

Author: mpage

Lib/test/test_peepholer.py

@@ -2637,6 +2637,32 @@ def test_format_simple(self):
         ]
         self.check(insts, expected)
 
+    def test_set_function_attribute(self):
+        # SET_FUNCTION_ATTRIBUTE leaves the function on the stack
+        insts = [
+            ("LOAD_CONST", 0, 1),
+            ("LOAD_FAST", 0, 2),
+            ("SET_FUNCTION_ATTRIBUTE", 2, 3),
+            ("STORE_FAST", 1, 4),
+            ("LOAD_CONST", 0, 5),
+            ("RETURN_VALUE", None, 6)
+        ]
+        self.cfg_optimization_test(insts, insts, consts=[None])
+
+        insts = [
+            ("LOAD_CONST", 0, 1),
+            ("LOAD_FAST", 0, 2),
+            ("SET_FUNCTION_ATTRIBUTE", 2, 3),
+            ("RETURN_VALUE", None, 4)
+        ]
+        expected = [
+            ("LOAD_CONST", 0, 1),
+            ("LOAD_FAST_BORROW", 0, 2),
+            ("SET_FUNCTION_ATTRIBUTE", 2, 3),
+            ("RETURN_VALUE", None, 4)
+        ]
+        self.cfg_optimization_test(insts, expected, consts=[None])
+
     def test_del_in_finally(self):
         # This loads `obj` onto the stack, executes `del obj`, then returns the
         # `obj` from the stack. See gh-133371 for more details.

Python/flowgraph.c

@@ -2944,6 +2944,16 @@ optimize_load_fast(cfg_builder *g)
                     break;
                 }
 
+                case SET_FUNCTION_ATTRIBUTE: {
+                    assert(_PyOpcode_num_popped(opcode, oparg) == 2);
+                    assert(_PyOpcode_num_pushed(opcode, oparg) == 1);
+                    ref func = ref_stack_pop(&refs);
+                    // Pop attr
+                    ref_stack_pop(&refs);
+                    PUSH_REF(func.instr, func.local);
+                    break;
+                }
+
                 // Opcodes that consume all of their inputs
                 default: {
                     int num_popped = _PyOpcode_num_popped(opcode, oparg);