RustPython · cthulahoops · Mar 8, 2019 · Mar 6, 2019 · Mar 8, 2019
diff --git 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
@@ -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
@@ -589,15 +589,15 @@ impl Compiler {
         args: &ast::Parameters,
         body: &[ast::LocatedStatement],
         decorator_list: &[ast::Expression],
-        _returns: &Option<ast::Expression>, // TODO: use type hint somehow..
+        returns: &Option<ast::Expression>, // 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
         let was_in_loop = self.in_loop;
         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
@@ -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)
             }