Support class __doc__

palaviv · palaviv · commit ab3cde77b4b2 · 2019-03-31T18:52:27.000+03:00
diff --git a/tests/snippets/class.py b/tests/snippets/class.py
@@ -37,8 +37,7 @@ def kungfu(x):
         assert x == 3
 
 
-# TODO:
-# assert Bar.__doc__ == " W00t "
+assert Bar.__doc__ == " W00t "
 
 bar = Bar(42)
 
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -769,6 +769,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
 pub fn builtin_build_class_(vm: &VirtualMachine, mut args: PyFuncArgs) -> PyResult {
     let function = args.shift();
     let name_arg = args.shift();
+    let doc = args.shift();
     let bases = args.args.clone();
     let mut metaclass = if let Some(metaclass) = args.get_optional_kwarg("metaclass") {
         PyClassRef::try_from_object(vm, metaclass)?
@@ -801,5 +802,6 @@ pub fn builtin_build_class_(vm: &VirtualMachine, mut args: PyFuncArgs) -> PyResu
         vec![name_arg, bases, namespace.into_object()],
     )?;
     cells.set_item(&vm.ctx, "__class__", class.clone());
+    vm.ctx.set_attr(&class, "__doc__", doc);
     Ok(class)
 }
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -706,13 +706,34 @@ impl Compiler {
             line_number,
             name.to_string(),
         ));
-        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)?;
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::None,
         });
         self.emit(Instruction::ReturnValue);
 
         let code = self.pop_code_object();
+
+        // function doc
+        self.emit(Instruction::LoadConst {
+            value: bytecode::Constant::String {
+                value: "".to_string(),
+            },
+        });
+
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::Code {
                 code: Box::new(code),
@@ -735,6 +756,11 @@ impl Compiler {
             },
         });
 
+        // class doc
+        self.emit(Instruction::LoadConst {
+            value: bytecode::Constant::String { value: doc_str },
+        });
+
         for base in bases {
             self.compile_expression(base)?;
         }
@@ -759,11 +785,11 @@ impl Compiler {
                 },
             });
             self.emit(Instruction::CallFunction {
-                typ: CallType::Keyword(2 + keywords.len() + bases.len()),
+                typ: CallType::Keyword(3 + keywords.len() + bases.len()),
             });
         } else {
             self.emit(Instruction::CallFunction {
-                typ: CallType::Positional(2 + bases.len()),
+                typ: CallType::Positional(3 + bases.len()),
             });
         }
 
@@ -1142,6 +1168,14 @@ impl Compiler {
                 self.compile_expression(body)?;
                 self.emit(Instruction::ReturnValue);
                 let code = self.pop_code_object();
+
+                // function doc
+                self.emit(Instruction::LoadConst {
+                    value: bytecode::Constant::String {
+                        value: "".to_string(),
+                    },
+                });
+
                 self.emit(Instruction::LoadConst {
                     value: bytecode::Constant::Code {
                         code: Box::new(code),
@@ -1428,6 +1462,13 @@ impl Compiler {
         // Fetch code for listcomp function:
         let code = self.pop_code_object();
 
+        // function doc
+        self.emit(Instruction::LoadConst {
+            value: bytecode::Constant::String {
+                value: "".to_string(),
+            },
+        });
+
         // List comprehension code:
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::Code {
