From b6dd92daf61eac069ebfe68ae5c7cb8d9e4b727e Mon Sep 17 00:00:00 2001
From: Windel Bouwman <windel@windel.nl>
Date: Wed, 7 Aug 2019 16:13:12 +0200
Subject: [PATCH] Properly fill the __annotations__ dictionary.

---
 src/compile.rs | 108 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 71 insertions(+), 37 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index 8d5c4c4..f7ac924 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -358,33 +358,7 @@ impl Compiler {
                 }
                 self.set_label(end_label);
             }
-            While { test, body, orelse } => {
-                let start_label = self.new_label();
-                let else_label = self.new_label();
-                let end_label = self.new_label();
-                self.emit(Instruction::SetupLoop {
-                    start: start_label,
-                    end: end_label,
-                });
-
-                self.set_label(start_label);
-
-                self.compile_test(test, None, Some(else_label), EvalContext::Statement)?;
-
-                let was_in_loop = self.in_loop;
-                self.in_loop = true;
-                self.compile_statements(body)?;
-                self.in_loop = was_in_loop;
-                self.emit(Instruction::Jump {
-                    target: start_label,
-                });
-                self.set_label(else_label);
-                self.emit(Instruction::PopBlock);
-                if let Some(orelse) = orelse {
-                    self.compile_statements(orelse)?;
-                }
-                self.set_label(end_label);
-            }
+            While { test, body, orelse } => self.compile_while(test, body, orelse)?,
             With {
                 is_async,
                 items,
@@ -557,16 +531,7 @@ impl Compiler {
                 target,
                 annotation,
                 value,
-            } => {
-                // Throw away annotation:
-                self.compile_expression(annotation)?;
-                self.emit(Instruction::Pop);
-
-                if let Some(value) = value {
-                    self.compile_expression(value)?;
-                    self.compile_store(target)?;
-                }
-            }
+            } => self.compile_annotated_assign(target, annotation, value)?,
             Delete { targets } => {
                 for target in targets {
                     self.compile_delete(target)?;
@@ -1015,6 +980,40 @@ impl Compiler {
         }
     }
 
+    fn compile_while(
+        &mut self,
+        test: &ast::Expression,
+        body: &[ast::Statement],
+        orelse: &Option<Vec<ast::Statement>>,
+    ) -> Result<(), CompileError> {
+        let start_label = self.new_label();
+        let else_label = self.new_label();
+        let end_label = self.new_label();
+        self.emit(Instruction::SetupLoop {
+            start: start_label,
+            end: end_label,
+        });
+
+        self.set_label(start_label);
+
+        self.compile_test(test, None, Some(else_label), EvalContext::Statement)?;
+
+        let was_in_loop = self.in_loop;
+        self.in_loop = true;
+        self.compile_statements(body)?;
+        self.in_loop = was_in_loop;
+        self.emit(Instruction::Jump {
+            target: start_label,
+        });
+        self.set_label(else_label);
+        self.emit(Instruction::PopBlock);
+        if let Some(orelse) = orelse {
+            self.compile_statements(orelse)?;
+        }
+        self.set_label(end_label);
+        Ok(())
+    }
+
     fn compile_for(
         &mut self,
         target: &ast::Expression,
@@ -1133,6 +1132,39 @@ impl Compiler {
         Ok(())
     }
 
+    fn compile_annotated_assign(
+        &mut self,
+        target: &ast::Expression,
+        annotation: &ast::Expression,
+        value: &Option<ast::Expression>,
+    ) -> Result<(), CompileError> {
+        if let Some(value) = value {
+            self.compile_expression(value)?;
+            self.compile_store(target)?;
+        }
+
+        // Compile annotation:
+        self.compile_expression(annotation)?;
+
+        if let ast::ExpressionType::Identifier { name } = &target.node {
+            // Store as dict entry in __annotations__ dict:
+            self.emit(Instruction::LoadName {
+                name: String::from("__annotations__"),
+                scope: bytecode::NameScope::Local,
+            });
+            self.emit(Instruction::LoadConst {
+                value: bytecode::Constant::String {
+                    value: name.to_string(),
+                },
+            });
+            self.emit(Instruction::StoreSubscript);
+        } else {
+            // Drop annotation if not assigned to simple identifier.
+            self.emit(Instruction::Pop);
+        }
+        Ok(())
+    }
+
     fn compile_store(&mut self, target: &ast::Expression) -> Result<(), CompileError> {
         match &target.node {
             ast::ExpressionType::Identifier { name } => {
@@ -1276,6 +1308,8 @@ impl Compiler {
 
     fn compile_expression(&mut self, expression: &ast::Expression) -> Result<(), CompileError> {
         trace!("Compiling {:?}", expression);
+        self.set_source_location(&expression.location);
+
         use ast::ExpressionType::*;
         match &expression.node {
             Call {