<code object>::replace() (#4207)

vm/src/builtins/code.rs

@@ -5,15 +5,39 @@
 use super::{PyStrRef, PyTupleRef, PyType, PyTypeRef};
 use crate::{
     builtins::PyStrInterned,
-    bytecode::{self, BorrowedConstant, Constant, ConstantBag},
+    bytecode::{self, BorrowedConstant, CodeFlags, Constant, ConstantBag},
     class::{PyClassImpl, StaticType},
     convert::ToPyObject,
-    function::FuncArgs,
+    function::{FuncArgs, OptionalArg},
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
 };
 use num_traits::Zero;
 use std::{borrow::Borrow, fmt, ops::Deref};
 
+#[derive(FromArgs)]
+pub struct ReplaceArgs {
+    #[pyarg(named, optional)]
+    co_posonlyargcount: OptionalArg<usize>,
+    #[pyarg(named, optional)]
+    co_argcount: OptionalArg<usize>,
+    #[pyarg(named, optional)]
+    co_kwonlyargcount: OptionalArg<usize>,
+    #[pyarg(named, optional)]
+    co_filename: OptionalArg<PyStrRef>,
+    #[pyarg(named, optional)]
+    co_firstlineno: OptionalArg<usize>,
+    #[pyarg(named, optional)]
+    co_consts: OptionalArg<Vec<PyObjectRef>>,
+    #[pyarg(named, optional)]
+    co_name: OptionalArg<PyStrRef>,
+    #[pyarg(named, optional)]
+    co_names: OptionalArg<Vec<PyObjectRef>>,
+    #[pyarg(named, optional)]
+    co_flags: OptionalArg<u16>,
+    #[pyarg(named, optional)]
+    co_varnames: OptionalArg<Vec<PyObjectRef>>,
+}
+
 #[derive(Clone)]
 pub struct Literal(PyObjectRef);
 
@@ -237,6 +261,93 @@ impl PyRef<PyCode> {
         let varnames = self.code.varnames.iter().map(|s| s.to_object()).collect();
         vm.ctx.new_tuple(varnames)
     }
+
+    #[pymethod]
+    pub fn replace(self, args: ReplaceArgs, vm: &VirtualMachine) -> PyResult<PyCode> {
+        let posonlyarg_count = match args.co_posonlyargcount {
+            OptionalArg::Present(posonlyarg_count) => posonlyarg_count,
+            OptionalArg::Missing => self.code.posonlyarg_count,
+        };
+
+        let arg_count = match args.co_argcount {
+            OptionalArg::Present(arg_count) => arg_count,
+            OptionalArg::Missing => self.code.arg_count,
+        };
+
+        let source_path = match args.co_filename {
+            OptionalArg::Present(source_path) => source_path,
+            OptionalArg::Missing => self.code.source_path.to_owned(),
+        };
+
+        let first_line_number = match args.co_firstlineno {
+            OptionalArg::Present(first_line_number) => first_line_number,
+            OptionalArg::Missing => self.code.first_line_number,
+        };
+
+        let kwonlyarg_count = match args.co_kwonlyargcount {
+            OptionalArg::Present(kwonlyarg_count) => kwonlyarg_count,
+            OptionalArg::Missing => self.code.kwonlyarg_count,
+        };
+
+        let constants = match args.co_consts {
+            OptionalArg::Present(constants) => constants,
+            OptionalArg::Missing => self.code.constants.iter().map(|x| x.0.clone()).collect(),
+        };
+
+        let obj_name = match args.co_name {
+            OptionalArg::Present(obj_name) => obj_name,
+            OptionalArg::Missing => self.code.obj_name.to_owned(),
+        };
+
+        let names = match args.co_names {
+            OptionalArg::Present(names) => names,
+            OptionalArg::Missing => self
+                .code
+                .names
+                .deref()
+                .iter()
+                .map(|name| name.to_pyobject(vm))
+                .collect(),
+        };
+
+        let flags = match args.co_flags {
+            OptionalArg::Present(flags) => flags,
+            OptionalArg::Missing => self.code.flags.bits(),
+        };
+
+        let varnames = match args.co_varnames {
+            OptionalArg::Present(varnames) => varnames,
+            OptionalArg::Missing => self.code.varnames.iter().map(|s| s.to_object()).collect(),
+        };
+
+        Ok(PyCode {
+            code: CodeObject {
+                flags: CodeFlags::from_bits_truncate(flags),
+                posonlyarg_count,
+                arg_count,
+                kwonlyarg_count,
+                source_path: source_path.as_object().as_interned_str(vm).unwrap(),
+                first_line_number,
+                obj_name: obj_name.as_object().as_interned_str(vm).unwrap(),
+
+                max_stackdepth: self.code.max_stackdepth,
+                instructions: self.code.instructions.clone(),
+                locations: self.code.locations.clone(),
+                constants: constants.into_iter().map(Literal).collect(),
+                names: names
+                    .into_iter()
+                    .map(|o| o.as_interned_str(vm).unwrap())
+                    .collect(),
+                varnames: varnames
+                    .into_iter()
+                    .map(|o| o.as_interned_str(vm).unwrap())
+                    .collect(),
+                cellvars: self.code.cellvars.clone(),
+                freevars: self.code.freevars.clone(),
+                cell2arg: self.code.cell2arg.clone(),
+            },
+        })
+    }
 }
 
 impl fmt::Display for PyCode {