Adds support for SETUP_ANNOTATION op and variable annotations

Tests have also been added. Fixes zrax#169
leonluffy · Oct 22, 2020 · 345035b · 345035b
1 parent d3a46f7
commit 345035b
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 3 deletions.
diff --git a/ASTNode.h b/ASTNode.h
@@ -16,6 +16,7 @@ class ASTNode {
         NODE_CONVERT, NODE_KEYWORD, NODE_RAISE, NODE_EXEC, NODE_BLOCK,
         NODE_COMPREHENSION, NODE_LOADBUILDCLASS, NODE_AWAITABLE,
         NODE_FORMATTEDVALUE, NODE_JOINEDSTR, NODE_CONST_MAP,
+        NODE_ANNOTATED_VAR,
 
         // Empty node types
         NODE_LOCALS,
@@ -669,4 +670,17 @@ class ASTJoinedStr : public ASTNode {
     value_t m_values;
 };
 
+class ASTAnnotatedVar : public ASTNode {
+public:
+    ASTAnnotatedVar(PycRef<ASTNode> name, PycRef<ASTNode> type)
+        : ASTNode(NODE_ANNOTATED_VAR), m_name(std::move(name)), m_type(std::move(type)) { }
+
+    const PycRef<ASTNode> name() const { return m_name; }
+    const PycRef<ASTNode> type() const { return m_type; }
+
+private:
+    PycRef<ASTNode> m_name;
+    PycRef<ASTNode> m_type;
+};
+
 #endif
diff --git a/ASTree.cpp b/ASTree.cpp
@@ -45,6 +45,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
     int unpack = 0;
     bool else_pop = false;
     bool need_try = false;
+    bool variable_annotations = false;
 
     while (!source.atEof()) {
 #if defined(BLOCK_DEBUG) || defined(STACK_DEBUG)
@@ -2117,10 +2118,36 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                     PycRef<ASTNode> src = stack.top();
                     stack.pop();
 
-                    if (dest.type() == ASTNode::NODE_MAP) {
-                        dest.cast<ASTMap>()->add(subscr, src);
+                    // If variable annotations are enabled, we'll need to check for them here.
+                    // Python handles a varaible annotation by setting:
+                    // __annotations__['var-name'] = type
+                    bool found_annotated_var = false;
+                    if (variable_annotations) {
+                        if (dest->type() == ASTNode::Type::NODE_NAME) {
+                            if (dest.cast<ASTName>()->name()->isEqual("__annotations__")) {
+                                found_annotated_var = true;
+                            }
+                        }
+                    }
+
+                    if (found_annotated_var) {
+                        // Annotations can be done alone or as part of an assignment.
+                        // In the case of an assignment, we'll see a NODE_STORE on the stack.
+                        if (curblock->nodes().back()->type() == ASTNode::Type::NODE_STORE) {
+                            // Replace the existing NODE_STORE with a new one that includes the annotation.
+                            PycRef<ASTStore> store = curblock->nodes().back().cast<ASTStore>();
+                            curblock->removeLast();
+                            curblock->append(new ASTStore(store->src(),
+                                                          new ASTAnnotatedVar(subscr, src)));
+                        } else {
+                            curblock->append(new ASTAnnotatedVar(subscr, src));
+                        }
                     } else {
-                        curblock->append(new ASTStore(src, new ASTSubscr(dest, subscr)));
+                        if (dest.type() == ASTNode::NODE_MAP) {
+                            dest.cast<ASTMap>()->add(subscr, src);
+                        } else {
+                            curblock->append(new ASTStore(src, new ASTSubscr(dest, subscr)));
+                        }
                     }
                 }
             }
@@ -2195,6 +2222,9 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 curblock->append(new ASTReturn(value, ASTReturn::YIELD));
             }
             break;
+        case Pyc::SETUP_ANNOTATIONS:
+            variable_annotations = true;
+            break;
         default:
             fprintf(stderr, "Unsupported opcode: %s\n", Pyc::OpcodeName(opcode & 0xFF));
             cleanBuild = false;
@@ -2973,6 +3003,17 @@ void print_src(PycRef<ASTNode> node, PycModule* mod)
                 fputc(')', pyc_output);
         }
         break;
+    case ASTNode::NODE_ANNOTATED_VAR:
+        {
+            PycRef<ASTAnnotatedVar> annotated_var = node.cast<ASTAnnotatedVar>();
+            PycRef<ASTObject> name = annotated_var->name().cast<ASTObject>();
+            PycRef<ASTNode> type = annotated_var->type();
+
+            fputs(name->object().cast<PycString>()->value(), pyc_output);
+            fputs(": ", pyc_output);
+            print_src(type, mod);
+        }
+        break;
     default:
         fprintf(pyc_output, "<NODE:%d>", node->type());
         fprintf(stderr, "Unsupported Node type: %d\n", node->type());

diff --git a/tests/compiled/variable_annotations.3.8.pyc b/tests/compiled/variable_annotations.3.8.pyc
diff --git a/tests/input/variable_annotations.py b/tests/input/variable_annotations.py
@@ -0,0 +1,8 @@
+a: int
+b
+c = 'no annotation'
+x: int = 10
+y: str = 'annotation'
+z: tuple = (1, 2, 3)
+confirm_subscr = {}
+confirm_subscr['test'] = 'works'
diff --git a/tests/tokenized/variable_annotations.txt b/tests/tokenized/variable_annotations.txt
@@ -0,0 +1,8 @@
+a : int <EOL>
+b <EOL>
+c = 'no annotation' <EOL>
+x : int = 10 <EOL>
+y : str = 'annotation' <EOL>
+z : tuple = ( 1 , 2 , 3 ) <EOL>
+confirm_subscr = { } <EOL>
+confirm_subscr [ 'test' ] = 'works' <EOL>