init

Author: windwhiterain

python/taichi/__init__.py

@@ -16,6 +16,8 @@
 # Issue#2223: Do not reorder, or we're busted with partially initialized module
 from taichi import aot  # isort:skip
 
+from taichi.lang.ast.hint import new, hide
+
 
 def __getattr__(attr):
     if attr == "cfg":

python/taichi/lang/ast/ast_transformer.py

@@ -16,6 +16,7 @@
 from taichi.lang._ndrange import _Ndrange, ndrange
 from taichi.lang.argpack import ArgPackType
 from taichi.lang.ast.ast_transformer_utils import Builder, LoopStatus, ReturnStatus
+from taichi.lang.ast import hint
 from taichi.lang.ast.symbol_resolver import ASTResolver
 from taichi.lang.exception import (
     TaichiIndexError,
@@ -105,18 +106,29 @@ def build_assign_annotated(ctx, target, value, is_static_assign, annotation):
                 f'Kernel argument "{target.id}" is immutable in the kernel. '
                 f"If you want to change its value, please create a new variable."
             )
-        anno = impl.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_ops.cast(value, anno)
-            var = impl.expr_init(var)
+        if annotation is hint.new:
+            var = impl.expr_init(value)
             ctx.create_variable(target.id, var)
+        elif annotation is hint.hide:
+            var = impl.expr_init(value)
+            ctx.hide_variable(target.id, var)
         else:
-            var = build_stmt(ctx, target)
-            if var.ptr.get_rvalue_type() != anno:
-                raise TaichiSyntaxError("Static assign cannot have type overloading")
-            var._assign(value)
+            anno = impl.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_ops.cast(value, anno)
+                var = impl.expr_init(var)
+                ctx.create_variable(target.id, var)
+            else:
+                var = build_stmt(ctx, target)
+                if var.ptr.get_rvalue_type() != anno:
+                    raise TaichiSyntaxError(
+                        "Static assign cannot have type overloading"
+                    )
+                var._assign(value)
         return var
 
     @staticmethod

python/taichi/lang/ast/ast_transformer_utils.py

@@ -240,10 +240,15 @@ def is_var_declared(self, name):
         return False
 
     def create_variable(self, name, var):
-        if name in self.current_scope():
-            raise TaichiSyntaxError("Recreating variables is not allowed")
+        # if name in self.current_scope():
+        #     raise TaichiSyntaxError("Recreating variables is not allowed")
         self.current_scope()[name] = var
 
+    def hide_variable(self, name, var):
+        for s in reversed(self.local_scopes):
+            if name in s:
+                s[name] = var
+
     def check_loop_var(self, loop_var):
         if self.is_var_declared(loop_var):
             raise TaichiSyntaxError(

python/taichi/lang/ast/hint.py

@@ -0,0 +1,5 @@
+from typing import Any
+
+
+type new = Any
+type hide = Any