diff --git a/tests/snippets/type_hints.py b/tests/snippets/type_hints.py index 8576e72352..0a36e36d6f 100644 --- a/tests/snippets/type_hints.py +++ b/tests/snippets/type_hints.py @@ -1,7 +1,7 @@ # See also: https://github.com/RustPython/RustPython/issues/587 -def curry(foo: int) -> float: - return foo * 3.1415926 * 2 +def curry(foo: int, bla=2) -> float: + return foo * 3.1415926 * bla assert curry(2) > 10 diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs index 59135ef64f..ff6ec873ce 100644 --- a/vm/src/bytecode.rs +++ b/vm/src/bytecode.rs @@ -31,6 +31,7 @@ pub struct CodeObject { bitflags! { pub struct FunctionOpArg: u8 { const HAS_DEFAULTS = 0x01; + const HAS_ANNOTATIONS = 0x04; } } diff --git a/vm/src/compile.rs b/vm/src/compile.rs index a5e01ab2f8..442aa3169c 100644 --- a/vm/src/compile.rs +++ b/vm/src/compile.rs @@ -589,7 +589,7 @@ impl Compiler { args: &ast::Parameters, body: &[ast::LocatedStatement], decorator_list: &[ast::Expression], - _returns: &Option, // TODO: use type hint somehow.. + returns: &Option, // TODO: use type hint somehow.. ) -> Result<(), CompileError> { // Create bytecode for this function: // remember to restore self.in_loop to the original after the function is compiled @@ -597,7 +597,7 @@ impl Compiler { let was_in_function_def = self.in_function_def; self.in_loop = false; self.in_function_def = true; - let flags = self.enter_function(name, args)?; + let mut flags = self.enter_function(name, args)?; self.compile_statements(body)?; // Emit None at end: @@ -608,6 +608,43 @@ impl Compiler { let code = self.pop_code_object(); self.prepare_decorators(decorator_list)?; + + // Prepare type annotations: + let mut num_annotations = 0; + + // Return annotation: + if let Some(annotation) = returns { + // key: + self.emit(Instruction::LoadConst { + value: bytecode::Constant::String { + value: "return".to_string(), + }, + }); + // value: + self.compile_expression(annotation)?; + num_annotations += 1; + } + + for arg in args.args.iter() { + if let Some(annotation) = &arg.annotation { + self.emit(Instruction::LoadConst { + value: bytecode::Constant::String { + value: arg.arg.to_string(), + }, + }); + self.compile_expression(&annotation)?; + num_annotations += 1; + } + } + + if num_annotations > 0 { + flags |= bytecode::FunctionOpArg::HAS_ANNOTATIONS; + self.emit(Instruction::BuildMap { + size: num_annotations, + unpack: false, + }); + } + self.emit(Instruction::LoadConst { value: bytecode::Constant::Code { code: Box::new(code), diff --git a/vm/src/frame.rs b/vm/src/frame.rs index 9777596070..11cd261692 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -447,15 +447,31 @@ impl Frame { bytecode::Instruction::MakeFunction { flags } => { let _qualified_name = self.pop_value(); let code_obj = self.pop_value(); + + let _annotations = if flags.contains(bytecode::FunctionOpArg::HAS_ANNOTATIONS) { + self.pop_value() + } else { + vm.new_dict() + }; + let defaults = if flags.contains(bytecode::FunctionOpArg::HAS_DEFAULTS) { self.pop_value() } else { vm.get_none() }; + // pop argc arguments // argument: name, args, globals let scope = self.scope.clone(); let obj = vm.ctx.new_function(code_obj, scope, defaults); + + let annotation_repr = vm.to_pystr(&_annotations)?; + + warn!( + "Type annotation must be stored in attribute! {:?}", + annotation_repr + ); + // TODO: use annotations with set_attr here! self.push_value(obj); Ok(None) }