Function support __doc__

palaviv · palaviv · commit f994f8660e48 · 2019-03-31T18:51:57.000+03:00
diff --git a/tests/snippets/function.py b/tests/snippets/function.py
@@ -1,7 +1,10 @@
 def foo():
+    """test"""
     return 42
 
 assert foo() == 42
+assert foo.__doc__ == "test"
+
 
 def my_func(a,):
     return a+2
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -597,7 +597,20 @@ impl Compiler {
         self.in_loop = false;
         self.in_function_def = true;
         let mut flags = self.enter_function(name, args)?;
-        self.compile_statements(body)?;
+
+        let doc = get_doc(body);
+        let (new_body, doc_str) = match doc {
+            Some(val) => {
+                if let Some((_, body_rest)) = body.split_first() {
+                    (body_rest, val)
+                } else {
+                    (body, "".to_string())
+                }
+            }
+            None => (body, "".to_string()),
+        };
+
+        self.compile_statements(new_body)?;
 
         // Emit None at end:
         self.emit(Instruction::LoadConst {
@@ -644,6 +657,10 @@ impl Compiler {
             });
         }
 
+        self.emit(Instruction::LoadConst {
+            value: bytecode::Constant::String { value: doc_str },
+        });
+
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::Code {
                 code: Box::new(code),
@@ -1511,6 +1528,19 @@ impl Compiler {
     }
 }
 
+fn get_doc(body: &[ast::LocatedStatement]) -> Option<String> {
+    if let Some(val) = body.get(0) {
+        if let ast::Statement::Expression { ref expression } = val.node {
+            if let ast::Expression::String { ref value } = expression {
+                if let ast::StringGroup::Constant { ref value } = value {
+                    return Some(value.to_string());
+                }
+            }
+        }
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
     use super::Compiler;
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -568,6 +568,8 @@ impl Frame {
                     .downcast()
                     .expect("Second to top value on the stack must be a code object");
 
+                let doc = self.pop_value();
+
                 let annotations = if flags.contains(bytecode::FunctionOpArg::HAS_ANNOTATIONS) {
                     self.pop_value()
                 } else {
@@ -586,6 +588,7 @@ impl Frame {
                 let obj = vm.ctx.new_function(code_obj, scope, defaults);
 
                 vm.ctx.set_attr(&obj, "__annotations__", annotations);
+                vm.ctx.set_attr(&obj, "__doc__", doc);
 
                 self.push_value(obj);
                 Ok(None)
