diff --git a/python/taichi/__init__.py b/python/taichi/__init__.py index 034052ed5bd73..0882b41aa3d93 100644 --- a/python/taichi/__init__.py +++ b/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": diff --git a/python/taichi/lang/ast/ast_transformer.py b/python/taichi/lang/ast/ast_transformer.py index ba7890d9b6b79..8c3cfedd588a1 100644 --- a/python/taichi/lang/ast/ast_transformer.py +++ b/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,25 @@ 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 diff --git a/python/taichi/lang/ast/ast_transformer_utils.py b/python/taichi/lang/ast/ast_transformer_utils.py index 4a10d5e3546c6..b3f96f32ea58b 100644 --- a/python/taichi/lang/ast/ast_transformer_utils.py +++ b/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( diff --git a/python/taichi/lang/ast/hint.py b/python/taichi/lang/ast/hint.py new file mode 100644 index 0000000000000..04731deaea97c --- /dev/null +++ b/python/taichi/lang/ast/hint.py @@ -0,0 +1,5 @@ +from typing import Any + + +type new = Any +type hide = Any