[Lang] Support annotated assignment

python/taichi/lang/ast/ast_transformer.py

@@ -18,6 +18,53 @@ def build_Name(ctx, node):
         node.ptr = ctx.get_var_by_name(node.id)
         return node
 
+    @staticmethod
+    def build_AnnAssign(ctx, node):
+
+        node.value = build_stmt(ctx, node.value)
+        node.target = build_stmt(ctx, node.target)
+        node.annotation = build_stmt(ctx, node.annotation)
+
+        is_static_assign = isinstance(
+            node.value, ast.Call) and ASTResolver.resolve_to(
+                node.value.func, ti.static, globals())
+
+        node.ptr = ASTTransformer.build_ann_assign_basic(
+            ctx, node.target, node.value.ptr, is_static_assign,
+            node.annotation.ptr)
+        return node
+
+    @staticmethod
+    def build_ann_assign_basic(ctx, target, value, is_static_assign,
+                               annotation):
+        """Build basic assginment like this: target : annotation = value.
+
+         Args:
+            ctx (ast_builder_utils.BuilderContext): The builder context.
+            target (ast.Name): A variable name. `target.id` holds the name as
+            a string.
+            annotation: A type we hope to assign to the target
+            value: A node representing the value.
+            is_static_assign: A boolean value indicating whether this is a static assignment
+        """
+        is_local = isinstance(target, ast.Name)
+        anno = ti.expr_init(annotation)
+        if is_static_assign:
+            raise TaichiSyntaxError(
+                "Static assign cannot be used on annotated assignment")
+        if is_local and not ctx.is_var_declared(target.id):
+            var = ti.cast(value, anno)
+            var = ti.expr_init(var)
+
+            ctx.create_variable(target.id, var)
+        else:
+            var = target.ptr
+            if var.ptr.get_ret_type() != anno:
+                raise TaichiSyntaxError(
+                    "Static assign cannot have type overloading")
+            var.assign(value)
+        return var
+
     @staticmethod
     def build_Assign(ctx, node):
         node.value = build_stmt(ctx, node.value)

tests/python/test_assign.py

@@ -57,3 +57,33 @@ def func_assign():
         assert a == 1
 
     func_assign()
+
+
+@ti.test(debug=True)
+def test_assign_ann():
+    @ti.kernel
+    def func_ann():
+        #need to introduce it
+        my_float = ti.f32
+        a: ti.i32 = 1
+        b: ti.f32 = a
+        d: my_float = 1
+        assert a == 1
+        assert b == 1.0
+        assert d == 1.0
+
+    func_ann()
+
+
+@ti.test()
+def test_assign_ann_over():
+    @ti.kernel
+    def func_ann_over():
+        my_int = ti.i32
+        d: my_int = 2
+        d: ti.f32 = 2.0
+
+    with pytest.raises(ti.lang.exception.TaichiCompilationError) as e:
+        func_ann_over()
+
+    assert e.type is ti.lang.exception.TaichiCompilationError