From 4c3c5c5ed1f0c1ad1fa68a293c344e903887de4b Mon Sep 17 00:00:00 2001
From: Noa <33094578+coolreader18@users.noreply.github.com>
Date: Mon, 30 Aug 2021 00:50:44 -0500
Subject: [PATCH] Add the __future__.annotations feature

---
 src/compile.rs     | 21 +++++++++++++++++----
 src/symboltable.rs | 25 ++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/compile.rs b/src/compile.rs
index 91781a2..2d877bf 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -67,6 +67,7 @@ struct Compiler {
     current_source_location: ast::Location,
     qualified_path: Vec<String>,
     done_with_future_stmts: bool,
+    future_annotations: bool,
     ctx: CompileContext,
     class_name: Option<String>,
     opts: CompileOpts,
@@ -196,6 +197,7 @@ impl Compiler {
             current_source_location: ast::Location::default(),
             qualified_path: Vec::new(),
             done_with_future_stmts: false,
+            future_annotations: false,
             ctx: CompileContext {
                 loop_data: None,
                 in_class: false,
@@ -1081,7 +1083,7 @@ impl Compiler {
                 value: "return".to_owned(),
             });
             // value:
-            self.compile_expression(annotation)?;
+            self.compile_annotation(annotation)?;
             num_annotations += 1;
         }
 
@@ -1096,7 +1098,7 @@ impl Compiler {
                 self.emit_constant(ConstantData::Str {
                     value: self.mangle(&arg.node.arg).into_owned(),
                 });
-                self.compile_expression(annotation)?;
+                self.compile_annotation(annotation)?;
                 num_annotations += 1;
             }
         }
@@ -1538,6 +1540,17 @@ impl Compiler {
         Ok(())
     }
 
+    fn compile_annotation(&mut self, annotation: &ast::Expr) -> CompileResult<()> {
+        if self.future_annotations {
+            self.emit_constant(ConstantData::Str {
+                value: annotation.to_string(),
+            });
+        } else {
+            self.compile_expression(annotation)?;
+        }
+        Ok(())
+    }
+
     fn compile_annotated_assign(
         &mut self,
         target: &ast::Expr,
@@ -1555,7 +1568,7 @@ impl Compiler {
         }
 
         // Compile annotation:
-        self.compile_expression(annotation)?;
+        self.compile_annotation(annotation)?;
 
         if let ast::ExprKind::Name { id, .. } = &target.node {
             // Store as dict entry in __annotations__ dict:
@@ -2423,7 +2436,7 @@ impl Compiler {
                 "nested_scopes" | "generators" | "division" | "absolute_import"
                 | "with_statement" | "print_function" | "unicode_literals" => {}
                 // "generator_stop" => {}
-                // "annotations" => {}
+                "annotations" => self.future_annotations = true,
                 other => {
                     return Err(self.error(CompileErrorType::InvalidFutureFeature(other.to_owned())))
                 }
diff --git a/src/symboltable.rs b/src/symboltable.rs
index c0fd9a2..602861f 100644
--- a/src/symboltable.rs
+++ b/src/symboltable.rs
@@ -537,6 +537,7 @@ struct SymbolTableBuilder {
     class_name: Option<String>,
     // Scope stack.
     tables: Vec<SymbolTable>,
+    future_annotations: bool,
 }
 
 /// Enum to indicate in what mode an expression
@@ -557,6 +558,7 @@ impl SymbolTableBuilder {
         let mut this = Self {
             class_name: None,
             tables: vec![],
+            future_annotations: false,
         };
         this.enter_scope("top", SymbolTableType::Module, 0);
         this
@@ -617,14 +619,31 @@ impl SymbolTableBuilder {
 
     fn scan_parameter_annotation(&mut self, parameter: &ast::Arg) -> SymbolTableResult {
         if let Some(annotation) = &parameter.node.annotation {
-            self.scan_expression(annotation, ExpressionContext::Load)?;
+            self.scan_annotation(annotation)?;
         }
         Ok(())
     }
 
+    fn scan_annotation(&mut self, annotation: &ast::Expr) -> SymbolTableResult {
+        if self.future_annotations {
+            Ok(())
+        } else {
+            self.scan_expression(annotation, ExpressionContext::Load)
+        }
+    }
+
     fn scan_statement(&mut self, statement: &ast::Stmt) -> SymbolTableResult {
         use ast::StmtKind::*;
         let location = statement.location;
+        if let ImportFrom { module, names, .. } = &statement.node {
+            if module.as_deref() == Some("__future__") {
+                for feature in names {
+                    if feature.name == "annotations" {
+                        self.future_annotations = true;
+                    }
+                }
+            }
+        }
         match &statement.node {
             Global { names } => {
                 for name in names {
@@ -655,7 +674,7 @@ impl SymbolTableBuilder {
                 self.scan_expressions(decorator_list, ExpressionContext::Load)?;
                 self.register_name(name, SymbolUsage::Assigned, location)?;
                 if let Some(expression) = returns {
-                    self.scan_expression(expression, ExpressionContext::Load)?;
+                    self.scan_annotation(expression)?;
                 }
                 self.enter_function(name, args, location.row())?;
                 self.scan_statements(body)?;
@@ -769,7 +788,7 @@ impl SymbolTableBuilder {
                         self.scan_expression(target, ExpressionContext::Store)?;
                     }
                 }
-                self.scan_expression(annotation, ExpressionContext::Load)?;
+                self.scan_annotation(annotation)?;
                 if let Some(value) = value {
                     self.scan_expression(value, ExpressionContext::Load)?;
                 }