Skip to content

Commit edb7934

Browse files
youknowoneyouknowone
committed
PyFunction constructor
1 parent 4c60303 commit edb7934

File tree

2 files changed

+106
-23
lines changed

2 files changed

+106
-23
lines changed

‎vm/src/builtins/function.rs‎

Lines changed: 104 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use super::{
88
#[cfg(feature = "jit")]
99
use crate::common::lock::OnceCell;
1010
use crate::common::lock::PyMutex;
11-
use crate::convert::ToPyObject;
11+
use crate::convert::{ToPyObject, TryFromObject};
1212
use crate::function::ArgMapping;
1313
use crate::object::{Traverse, TraverseFn};
1414
use crate::{
@@ -31,6 +31,7 @@ use rustpython_jit::CompiledCode;
3131
pub struct PyFunction {
3232
code: PyRef<PyCode>,
3333
globals: PyDictRef,
34+
builtins: PyObjectRef,
3435
closure: Option<PyTupleTyped<PyCellRef>>,
3536
defaults_and_kwdefaults: PyMutex<(Option<PyTupleRef>, Option<PyDictRef>)>,
3637
name: PyMutex<PyStrRef>,
@@ -53,6 +54,7 @@ unsafe impl Traverse for PyFunction {
5354

5455
impl PyFunction {
5556
#[allow(clippy::too_many_arguments)]
57+
#[inline]
5658
pub(crate) fn new(
5759
code: PyRef<PyCode>,
5860
globals: PyDictRef,
@@ -62,24 +64,42 @@ impl PyFunction {
6264
qualname: PyStrRef,
6365
type_params: PyTupleRef,
6466
annotations: PyDictRef,
65-
module: PyObjectRef,
6667
doc: PyObjectRef,
67-
) -> Self {
68+
vm: &VirtualMachine,
69+
) -> PyResult<Self> {
6870
let name = PyMutex::new(code.obj_name.to_owned());
69-
PyFunction {
71+
let module = vm.unwrap_or_none(globals.get_item_opt(identifier!(vm, __name__), vm)?);
72+
let builtins = globals.get_item("__builtins__", vm).unwrap_or_else(|_| {
73+
// If not in globals, inherit from current execution context
74+
if let Some(frame) = vm.current_frame() {
75+
frame.builtins.clone().into()
76+
} else {
77+
vm.builtins.clone().into()
78+
}
79+
});
80+
81+
let func = PyFunction {
7082
code,
7183
globals,
84+
builtins,
7285
closure,
7386
defaults_and_kwdefaults: PyMutex::new((defaults, kw_only_defaults)),
7487
name,
7588
qualname: PyMutex::new(qualname),
7689
type_params: PyMutex::new(type_params),
77-
#[cfg(feature = "jit")]
78-
jitted_code: OnceCell::new(),
7990
annotations: PyMutex::new(annotations),
8091
module: PyMutex::new(module),
8192
doc: PyMutex::new(doc),
82-
}
93+
#[cfg(feature = "jit")]
94+
jitted_code: OnceCell::new(),
95+
};
96+
97+
// let name = qualname.as_str().split('.').next_back().unwrap();
98+
// func.set_attr(identifier!(vm, __name__), vm.new_pyobj(name), vm)?;
99+
// func.set_attr(identifier!(vm, __qualname__), qualname, vm)?;
100+
// func.set_attr(identifier!(vm, __doc__), doc, vm)?;
101+
102+
Ok(func)
83103
}
84104

85105
fn fill_locals_from_args(
@@ -362,7 +382,7 @@ impl PyPayload for PyFunction {
362382
}
363383

364384
#[pyclass(
365-
with(GetDescriptor, Callable, Representable),
385+
with(GetDescriptor, Callable, Representable, Constructor),
366386
flags(HAS_DICT, METHOD_DESCRIPTOR)
367387
)]
368388
impl PyFunction {
@@ -409,12 +429,7 @@ impl PyFunction {
409429
#[pymember(magic)]
410430
fn builtins(vm: &VirtualMachine, zelf: PyObjectRef) -> PyResult {
411431
let zelf = Self::_as_pyref(&zelf, vm)?;
412-
// Get __builtins__ from the function's globals dict
413-
let builtins = zelf
414-
.globals
415-
.get_item("__builtins__", vm)
416-
.unwrap_or_else(|_| vm.builtins.clone().into());
417-
Ok(builtins)
432+
Ok(zelf.builtins.clone())
418433
}
419434

420435
#[pygetset(magic)]
@@ -561,6 +576,81 @@ impl Representable for PyFunction {
561576
}
562577
}
563578

579+
#[derive(FromArgs)]
580+
pub struct PyFunctionNewArgs {
581+
#[pyarg(positional)]
582+
code: PyRef<PyCode>,
583+
#[pyarg(positional)]
584+
globals: PyDictRef,
585+
#[pyarg(any, optional)]
586+
name: OptionalArg<PyStrRef>,
587+
#[pyarg(any, optional)]
588+
defaults: OptionalArg<PyTupleRef>,
589+
#[pyarg(any, optional)]
590+
closure: OptionalArg<PyTupleRef>,
591+
#[pyarg(any, optional)]
592+
kwdefaults: OptionalArg<PyDictRef>,
593+
}
594+
595+
impl Constructor for PyFunction {
596+
type Args = PyFunctionNewArgs;
597+
598+
fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
599+
// Handle closure - must be a tuple of cells
600+
let closure = if let Some(closure_tuple) = args.closure.into_option() {
601+
// Check that closure length matches code's free variables
602+
if closure_tuple.len() != args.code.freevars.len() {
603+
return Err(vm.new_value_error(format!(
604+
"{} requires closure of length {}, not {}",
605+
args.code.obj_name,
606+
args.code.freevars.len(),
607+
closure_tuple.len()
608+
)));
609+
}
610+
611+
// Validate that all items are cells and create typed tuple
612+
let typed_closure =
613+
PyTupleTyped::<PyCellRef>::try_from_object(vm, closure_tuple.into())?;
614+
Some(typed_closure)
615+
} else if !args.code.freevars.is_empty() {
616+
return Err(vm.new_type_error("arg 5 (closure) must be tuple".to_owned()));
617+
} else {
618+
None
619+
};
620+
621+
// Get function name - use provided name or default to code object name
622+
let name = args
623+
.name
624+
.into_option()
625+
.unwrap_or_else(|| PyStr::from(args.code.obj_name.as_str()).into_ref(&vm.ctx));
626+
627+
// Get qualname - for now just use the name
628+
let qualname = name.clone();
629+
630+
// Create empty type_params and annotations
631+
let type_params = vm.ctx.new_tuple(vec![]);
632+
let annotations = vm.ctx.new_dict();
633+
634+
// Get doc from code object - for now just use None
635+
let doc = vm.ctx.none();
636+
637+
let func = PyFunction::new(
638+
args.code,
639+
args.globals,
640+
closure,
641+
args.defaults.into_option(),
642+
args.kwdefaults.into_option(),
643+
qualname,
644+
type_params,
645+
annotations,
646+
doc,
647+
vm,
648+
)?;
649+
650+
func.into_ref_with_type(vm, cls).map(Into::into)
651+
}
652+
}
653+
564654
#[pyclass(module = false, name = "method", traverse)]
565655
#[derive(Debug)]
566656
pub struct PyBoundMethod {

‎vm/src/frame.rs‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1905,8 +1905,6 @@ impl ExecutingFrame<'_> {
19051905
None
19061906
};
19071907

1908-
let module = vm.unwrap_or_none(self.globals.get_item_opt(identifier!(vm, __name__), vm)?);
1909-
19101908
// pop argc arguments
19111909
// argument: name, args, globals
19121910
// let scope = self.scope.clone();
@@ -1919,16 +1917,11 @@ impl ExecutingFrame<'_> {
19191917
qualified_name.clone(),
19201918
type_params,
19211919
annotations.downcast().unwrap(),
1922-
module,
19231920
vm.ctx.none(),
1924-
)
1921+
vm,
1922+
)?
19251923
.into_pyobject(vm);
19261924

1927-
let name = qualified_name.as_str().split('.').next_back().unwrap();
1928-
func_obj.set_attr(identifier!(vm, __name__), vm.new_pyobj(name), vm)?;
1929-
func_obj.set_attr(identifier!(vm, __qualname__), qualified_name, vm)?;
1930-
func_obj.set_attr(identifier!(vm, __doc__), vm.ctx.none(), vm)?;
1931-
19321925
self.push_value(func_obj);
19331926
Ok(None)
19341927
}

0 commit comments

Comments
 (0)