From 16a896d51311ec842b7c48314d675c2d33d9b095 Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 28 Sep 2020 23:52:59 -0500
Subject: [PATCH] Fix accessing __doc__ while inside a class definition

---
 src/compile.rs     | 22 +++++++++++++---------
 src/symboltable.rs |  1 +
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index 2c04ecc..f92eadf 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -989,7 +989,13 @@ impl<O: OutputStream> Compiler<O> {
 
         // Turn code object into function object:
         self.emit(Instruction::MakeFunction);
-        self.store_docstring(doc_str);
+
+        self.emit(Instruction::Duplicate);
+        self.load_docstring(doc_str);
+        self.emit(Instruction::Rotate { amount: 2 });
+        self.emit(Instruction::StoreAttr {
+            name: "__doc__".to_owned(),
+        });
         self.apply_decorators(decorator_list);
 
         self.store_name(name);
@@ -1109,6 +1115,11 @@ impl<O: OutputStream> Compiler<O> {
             name: "__qualname__".to_owned(),
             scope: bytecode::NameScope::Free,
         });
+        self.load_docstring(doc_str);
+        self.emit(Instruction::StoreName {
+            name: "__doc__".to_owned(),
+            scope: bytecode::NameScope::Free,
+        });
         // setup annotations
         if self.find_ann(body) {
             self.emit(Instruction::SetupAnnotation);
@@ -1175,7 +1186,6 @@ impl<O: OutputStream> Compiler<O> {
             });
         }
 
-        self.store_docstring(doc_str);
         self.apply_decorators(decorator_list);
 
         self.store_name(name);
@@ -1184,10 +1194,9 @@ impl<O: OutputStream> Compiler<O> {
         Ok(())
     }
 
-    fn store_docstring(&mut self, doc_str: Option<String>) {
+    fn load_docstring(&mut self, doc_str: Option<String>) {
         // TODO: __doc__ must be default None and no bytecodes unless it is Some
         // Duplicate top of stack (the function or class object)
-        self.emit(Instruction::Duplicate);
 
         // Doc string value:
         self.emit(Instruction::LoadConst {
@@ -1196,11 +1205,6 @@ impl<O: OutputStream> Compiler<O> {
                 None => bytecode::Constant::None, // set docstring None if not declared
             },
         });
-
-        self.emit(Instruction::Rotate { amount: 2 });
-        self.emit(Instruction::StoreAttr {
-            name: "__doc__".to_owned(),
-        });
     }
 
     fn compile_while(
diff --git a/src/symboltable.rs b/src/symboltable.rs
index 4d759f9..0ad497c 100644
--- a/src/symboltable.rs
+++ b/src/symboltable.rs
@@ -531,6 +531,7 @@ impl SymbolTableBuilder {
                 self.enter_scope(name, SymbolTableType::Class, statement.location.row());
                 self.register_name("__module__", SymbolUsage::Assigned)?;
                 self.register_name("__qualname__", SymbolUsage::Assigned)?;
+                self.register_name("__doc__", SymbolUsage::Assigned)?;
                 self.scan_statements(body)?;
                 self.leave_scope();
                 self.scan_expressions(bases, &ExpressionContext::Load)?;