From a3b7a837417aeebee2bc892a6a2bc9403368a4cc Mon Sep 17 00:00:00 2001
From: lynskylate <lynskylate@gmail.com>
Date: Tue, 7 Jul 2020 06:00:50 +0800
Subject: [PATCH] Add annotation setup

---
 src/compile.rs | 70 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/src/compile.rs b/src/compile.rs
index 7303e43..ceb3743 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -984,6 +984,64 @@ impl<O: OutputStream> Compiler<O> {
         Ok(())
     }
 
+    fn option_stmt_to_bool(&mut self, stmts: &Option<ast::Suite>) -> bool {
+        match &stmts {
+            Some(stmts) => self.find_ann(stmts),
+            None => false,
+        }
+    }
+
+    fn find_ann(&mut self, body: &[ast::Statement]) -> bool {
+        use ast::StatementType::*;
+
+        for statement in body {
+            let res = match &statement.node {
+                AnnAssign {
+                    target: _,
+                    annotation: _,
+                    value: _,
+                } => true,
+                For {
+                    is_async: _,
+                    target: _,
+                    iter: _,
+                    body,
+                    orelse,
+                } => self.find_ann(body) || self.option_stmt_to_bool(orelse),
+                If {
+                    test: _,
+                    body,
+                    orelse,
+                } => self.find_ann(body) || self.option_stmt_to_bool(orelse),
+                While {
+                    test: _,
+                    body,
+                    orelse,
+                } => self.find_ann(body) || self.option_stmt_to_bool(orelse),
+                With {
+                    is_async: _,
+                    items: _,
+                    body,
+                } => self.find_ann(body),
+                Try {
+                    body,
+                    handlers: _,
+                    orelse,
+                    finalbody,
+                } => {
+                    self.find_ann(body)
+                        || self.option_stmt_to_bool(orelse)
+                        || self.option_stmt_to_bool(finalbody)
+                }
+                _ => false,
+            };
+            if res {
+                return true;
+            }
+        }
+        return false;
+    }
+
     fn compile_class_def(
         &mut self,
         name: &str,
@@ -1037,6 +1095,18 @@ impl<O: OutputStream> Compiler<O> {
             name: "__qualname__".to_owned(),
             scope: bytecode::NameScope::Free,
         });
+        // setup annotations
+        if self.find_ann(body) {
+            self.emit(Instruction::BuildMap {
+                size: 0,
+                unpack: false,
+                for_call: false,
+            });
+            self.emit(Instruction::StoreName {
+                name: "__annotations__".to_owned(),
+                scope: bytecode::NameScope::Free,
+            });
+        }
         self.compile_statements(new_body)?;
         self.emit(Instruction::LoadConst {
             value: bytecode::Constant::None,