From c9679bd7c36df9f43bd1d6a0869a0c178239e666 Mon Sep 17 00:00:00 2001
From: archibate <17721388340@163.com>
Date: Thu, 27 Feb 2020 14:22:21 +0800
Subject: [PATCH] add TI_PRINT_PREPROCESSED

fix TI_PRINT_PROCESSED
---
 python/taichi/lang/__init__.py    | 5 +++++
 python/taichi/lang/transformer.py | 5 ++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py
index 27d9944196bac0..bd3c9b3d1f404c 100644
--- a/python/taichi/lang/__init__.py
+++ b/python/taichi/lang/__init__.py
@@ -76,6 +76,11 @@ def init(default_fp=None, default_ip=None, print_preprocessed=None, debug=None,
     elif dfl_ip is not None:
       raise ValueError(f'Unrecognized TI_DEFAULT_IP: {dfl_ip}, should be 32 or 64')
 
+  if print_preprocessed is None: # won't override
+    print_preprocessed = os.environ.get("TI_PRINT_PREPROCESSED")
+    if print_preprocessed is not None:
+      print_preprocessed = bool(int(print_preprocessed))
+
   if default_fp is not None:
     ti.get_runtime().set_default_fp(default_fp)
   if default_ip is not None:
diff --git a/python/taichi/lang/transformer.py b/python/taichi/lang/transformer.py
index 368eca1143e389..86f1ffc2d1a6f8 100644
--- a/python/taichi/lang/transformer.py
+++ b/python/taichi/lang/transformer.py
@@ -216,8 +216,11 @@ def visit_block(self, list_stmt):
     for i, l in enumerate(list_stmt):
       list_stmt[i] = self.visit(l)
 
+  def visit_Return(self, node):
+    print('visited return!')
+
   def visit_If(self, node):
-    self.generic_visit(node, ['body', 'orelse'])
+    self.generic_visit(node)
 
     is_static_if = isinstance(node.test, ast.Call) and isinstance(
         node.test.func, ast.Attribute)