mindspore 2.1.0__cp37-cp37m-manylinux1_x86_64.whl → 2.2.10__cp37-cp37m-manylinux1_x86_64.whl

mindspore/rewrite/parsers/attribute_parser.py

@@ -15,13 +15,11 @@
 """Parse ast.Assign in construct function to node of SymbolTree."""
 import ast
 
-from mindspore.rewrite.parser import Parser
+from mindspore.rewrite.parsers.parser import Parser
 from mindspore.rewrite.symbol_tree import SymbolTree
-from mindspore.rewrite.parser_register import ParserRegister
-
-from mindspore.rewrite.parser_register import reg_parser
+from mindspore.rewrite.parsers.parser_register import ParserRegister, reg_parser
 from ..common import error_str
-
+from ..node.node_manager import NodeManager
 
 class AttributeParser(Parser):
     """Parse ast.Attribute in construct function to node of SymbolTree."""
@@ -30,13 +28,14 @@ class AttributeParser(Parser):
         """Parse target type."""
         return ast.Attribute
 
-    def process(self, stree: SymbolTree, node: ast.Attribute):
+    def process(self, stree: SymbolTree, node: ast.Attribute, node_manager: NodeManager):
         """
         Parse ast.Attribute node.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([ast.Attribute]): An ast.Attribute node.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Returns:
             The value of node.
@@ -47,8 +46,11 @@ class AttributeParser(Parser):
         if not isinstance(node, ast.Attribute):
             raise TypeError(error_str(f"Attribute parser only supports parsing ast.Attribute type nodes, but got "
                                       f"'{type(node).__name__}'", father_node=node))
+        if not isinstance(node.value, (ast.Name, ast.Attribute)):
+            raise RuntimeError(error_str(f"Attribute parser only supports (ast.Attribute, ast.Name) as value of "
+                                         f"ast.Attribute, but got '{type(node).__name__}'", father_node=node))
         parser = ParserRegister.instance().get_parser(type(node.value))
-        value = parser.process(stree, node.value)
+        value = parser.process(stree, node.value, node_manager)
 
         return ".".join([value, node.attr])
 

mindspore/rewrite/parsers/class_def_parser.py

@@ -13,17 +13,19 @@
 # limitations under the License.
 # ============================================================================
 """Parse ast.ClassDef which is subclass of Cell to SymbolTree."""
-import sys
-import ast
 import inspect
+from typing import Union, Dict
+import ast
 from mindspore import log as logger
 from mindspore.nn import Cell
 from mindspore._extends.parse.namespace import CellNamespace
 from ..symbol_tree import SymbolTree
-from ..parser import Parser
-from ..parser_register import ParserRegister, reg_parser
+from .parser import Parser
+from .parser_register import ParserRegister, reg_parser
 from ..ast_helpers import AstReplacer
 from ..common import error_str
+from ..parsers.module_parser import ModuleParser
+from ..node.node_manager import NodeManager
 
 
 class AstScopeChecker:
@@ -106,25 +108,58 @@ class AstScopeChecker:
 class ClassDefParser(Parser):
     """Parse ast.ClassDef which is subclass of Cell to SymbolTree."""
 
+    # a denied_function_decorator_list which is registered by user
+    denied_function_decorator_list = []
+    # Entry function of the forward computation process
+    entry_function = "construct"
+
     def __init__(self):
         """Constructor"""
         super(ClassDefParser, self).__init__()
         self._cell_namespace = CellNamespace('mindspore.nn')
 
     @staticmethod
-    def _is_super_expr(expr: ast.AST) -> bool:
-        """Check whether ast node is super().__init__()"""
-        if not isinstance(expr, ast.Expr):
-            return False
-        if not isinstance(expr.value, ast.Call):
-            return False
-        if not isinstance(expr.value.func, ast.Attribute):
-            return False
-        if expr.value.func.attr != "__init__" or not isinstance(expr.value.func.value, ast.Call):
-            return False
-        if not isinstance(expr.value.func.value.func, ast.Name) or expr.value.func.value.func.id != "super":
-            return False
-        return True
+    def _process_init_func_ast(init_ast: ast.FunctionDef, class_name: str, is_father_class: bool,
+                               father_classes: dict):
+        """Process init func"""
+        ClassDefParser._modify_arguments_of_init_func(init_ast)
+        new_bodies = ClassDefParser._create_bodys_of_init_func(class_name, is_father_class, father_classes)
+        init_ast.body = new_bodies
+
+    @staticmethod
+    def _create_bodys_of_init_func(class_name: str, is_father_class: bool, father_classes: dict):
+        """Modify bodys of init func."""
+        new_bodies = []
+        # update father class init in new class
+        father_class_init_bodies = ClassDefParser._father_class_init_process(father_classes, is_father_class)
+        new_bodies.extend(father_class_init_bodies)
+        # copy variables into new class
+        if is_father_class:
+            ast_copy_attr = ast.parse(
+                "for key, value in obj.__dict__.items():\n"
+                "    if not key.startswith('__'):\n"
+                f"        setattr({class_name}, key, value)").body[0]
+            new_bodies.append(ast_copy_attr)
+        else:
+            ast_copy_attr = ast.parse(
+                "for key, value in obj.__dict__.items(): setattr(self, key, value)").body[0]
+            new_bodies.append(ast_copy_attr)
+        return new_bodies
+
+    @staticmethod
+    def _father_class_init_process(father_classes: dict, is_father_class: bool) -> [ast.AST]:
+        """Add ast bodies of code: father_class.__init__(...)"""
+        father_class_init_bodies = []
+        for idx, father_class in father_classes.items():
+            if father_class == "Cell":
+                father_class_init_code = "super().__init__()"
+            elif is_father_class:
+                father_class_init_code = f"{father_class}.__init__(self, obj.__bases__[{idx}])"
+            else:
+                father_class_init_code = f"{father_class}.__init__(self, obj.__class__.__bases__[{idx}])"
+            father_class_init_ast = ast.parse(father_class_init_code).body[0]
+            father_class_init_bodies.append(father_class_init_ast)
+        return father_class_init_bodies
 
     @staticmethod
     def _modify_arguments_of_init_func(ast_init_fn: ast.FunctionDef):
@@ -135,127 +170,153 @@ class ClassDefParser(Parser):
                                          kw_defaults=[], defaults=[], vararg=None, kwarg=None)
         ast.fix_missing_locations(ast_init_fn)
 
+    @staticmethod
+    def get_ast_name(ast_node: Union[ast.Name, ast.Attribute]) -> str:
+        """Get ast id name"""
+        if isinstance(ast_node, ast.Name):
+            return ast_node.id
+        if isinstance(ast_node, ast.Attribute):
+            return ast_node.attr
+        return ""
+
+    @staticmethod
+    def _process_class_variables(stree: SymbolTree, function_defs: list):
+        """Process class variables of class, only used in child class."""
+        init_func_ast = stree.get_init_func_ast()
+        for key, value in stree.get_origin_network().__class__.__dict__.items():
+            if key.startswith('__'):
+                # ignore inner functions
+                continue
+            if callable(value) and key in function_defs:
+                # ignore functions defined by self
+                continue
+            assign_code = f"self.__class__.{key} = obj.__class__.{key}"
+            assign_ast = ast.parse(assign_code).body[0]
+            init_func_ast.body.append(assign_ast)
+
+    @staticmethod
+    def _need_add_init_func(cls_ast: ast.ClassDef) -> bool:
+        """If the class don't have init func, we need to add an init func"""
+        for body in cls_ast.body:
+            if isinstance(body, ast.FunctionDef) and body.name == '__init__':
+                return False
+        return True
+
+    @staticmethod
+    def _add_init_func(cls_ast: ast.ClassDef):
+        """Add init func with super().__init__()"""
+        init_func_ast = ast.parse("def __init__(self): super().__init__()").body[0]
+        cls_ast.body.insert(0, init_func_ast)
+        ast.fix_missing_locations(cls_ast)
+
+    @staticmethod
+    def _process_father_classes(stree, node: ast.ClassDef, cur_class_def: type) -> list:
+        """Process father class."""
+        father_classes: Dict[int, str] = {}
+        for idx, base in enumerate(node.bases):
+            father_class_name = ClassDefParser.get_ast_name(base)
+            if not father_class_name:
+                continue
+            father_classes[idx] = father_class_name
+            if father_class_name == "Cell":
+                continue
+            father_class_def = cur_class_def.__bases__[idx]
+            ClassDefParser._process_one_father_class(stree, father_class_def, father_class_name)
+            node.bases[idx] = ast.Name(id=father_class_name, ctx=ast.Load())
+        return father_classes
+
+    @staticmethod
+    def _process_one_father_class(stree: SymbolTree, father_class_def: type, father_class_name: str):
+        """Process one father class"""
+        # save father class's file path and imports into symbol tree
+        net_path = inspect.getfile(father_class_def)
+        ModuleParser.save_file_path_to_sys(stree, 0, net_path)
+        ModuleParser.save_imports_from_file(stree, net_path)
+        # get father class's ast
+        source_code = inspect.getsource(father_class_def)
+        father_class_ast: ast.ClassDef = ast.parse(source_code).body[0]
+        # process father class's father classes
+        father_classes = ClassDefParser._process_father_classes(stree, father_class_ast, father_class_def)
+        # process father class's __init__ function
+        if ClassDefParser._need_add_init_func(father_class_ast):
+            ClassDefParser._add_init_func(father_class_ast)
+        for body in father_class_ast.body[:]:
+            if isinstance(body, ast.FunctionDef) and body.name == "__init__":
+                # Add function decorator
+                ClassDefParser._func_decorator_process(body)
+                ClassDefParser._process_init_func_ast(body, father_class_name, True, father_classes)
+            else:
+                # Remove other codes, which are copied in __init__ function.
+                father_class_ast.body.remove(body)
+        # save father class's ast into symbol tree
+        stree.get_father_class_ast().append(father_class_ast)
+
+    @staticmethod
+    def _func_decorator_process(node: ast.FunctionDef):
+        """
+        User should set the denied function decorators,
+        because the symbol_tree cant pass the correct parameters to decorators but the instance "obj".
+        """
+        for decorator in node.decorator_list[:]:
+            decorator_name = ""
+            if isinstance(decorator, ast.Call):
+                func = decorator.func
+                if isinstance(func, ast.Name):
+                    decorator_name = func.id
+            elif isinstance(decorator, ast.Name):
+                decorator_name = decorator.id
+            if decorator_name in ClassDefParser.denied_function_decorator_list:
+                node.decorator_list.remove(decorator)
+
     def target(self):
         """Parse target type"""
         return ast.ClassDef
 
-    def process(self, stree: SymbolTree, node: ast.ClassDef):
+    def process(self, stree: SymbolTree, node: ast.ClassDef, node_manager: NodeManager):
         """
-        Parse init and construct in ast.ClassDef.
+        Parse init and entry function(default: construct) in ast.ClassDef.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([ast.ClassDef]): An ast.ClassDef node.
+            node_manager (NodeManager): NodeManager those asts belong to.
         """
+        # Update network's class name from xxx to xxxOpt in ast
         replacer = AstReplacer(node)
         replacer.replace_all(stree.get_ori_cls_name(), stree.get_opt_cls_name())
 
+        # process network's father classes
         stree.set_class_ast(node)
-        has_father_class = self._handle_father_class(stree, node)
+        cur_class_def = type(stree.get_origin_network())
+        father_classes = ClassDefParser._process_father_classes(stree, node, cur_class_def)
 
-        if self._need_add_init_func(stree, node):
-            self._add_init_func(node)
+        # add __init__ function to network if necessary
+        if isinstance(stree.get_origin_network(), Cell) and ClassDefParser._need_add_init_func(node):
+            ClassDefParser._add_init_func(node)
 
-        for body in node.body:
+        # save function defs in ast node to filter function class variables.
+        function_defs = []
+        for body in node.body[:]:
             if isinstance(body, ast.FunctionDef):
+                function_defs.append(body.name)
+                ClassDefParser._func_decorator_process(body)
                 if body.name == "__init__":
-                    self._process_init_func_ast(body, has_father_class)
                     stree.set_init_func_ast(body)
-                elif body.name == "construct":
+                    ClassDefParser._process_init_func_ast(body, stree.get_opt_cls_name(), False, father_classes)
+                elif body.name == ClassDefParser.entry_function:
+                    stree.set_ast_root(body)
                     parser: Parser = ParserRegister.instance().get_parser(ast.FunctionDef)
-                    parser.process(stree, body)
+                    parser.process(stree, body, stree)
                 else:
                     logger.info(
                         "Ignoring ast.FunctionDef in ast.ClassDef except __init__ and construct function: %s",
                         body.name)
+            elif isinstance(body, (ast.Assign, ast.If, ast.IfExp)):
+                # Remove class variables, which are copied in __init__ function.
+                node.body.remove(body)
             else:
                 logger.info("Ignoring unsupported node(%s) in ast.ClassDef.", type(body).__name__)
-
-    def _is_subtree_field(self, ori_net, field) -> bool:
-        op = getattr(ori_net, field)
-        return not type(op).__name__ in self._cell_namespace
-
-    def _process_init_func_ast(self, init_ast: ast.FunctionDef, has_father_class: bool):
-        """Process init func"""
-        ClassDefParser._modify_arguments_of_init_func(init_ast)
-        new_bodies = self._replace_ori_field_of_init_func(init_ast.body, has_father_class)
-        init_ast.body = new_bodies
-
-    def _need_add_init_func(self, stree: SymbolTree, cls_ast: ast.ClassDef) -> bool:
-        """If class is child class of nn.Cell but not have init func, then we need to add init func"""
-        if not isinstance(stree.get_origin_network(), Cell):
-            return False
-        for body in cls_ast.body:
-            if isinstance(body, ast.FunctionDef) and body.name == '__init__':
-                return False
-        return True
-
-    def _add_init_func(self, cls_ast: ast.ClassDef):
-        """Add init func with super().__init__()"""
-        init_func_ast = ast.parse("def __init__(self): super().__init__()").body[0]
-        cls_ast.body.insert(0, init_func_ast)
-        ast.fix_missing_locations(cls_ast)
-
-    def _replace_ori_field_of_init_func(self, bodies: [], has_father_class: bool):
-        """
-        Replace original field in init func to self.XX = getattr(self._handler, "XX").
-        Only keep following two kinds of ast nodes in bodies right now:
-            1. Ast.If and test is self.XX.
-            2. Ast.Assign and target is self.XX.
-
-        Args:
-            bodies ([]): bodied of init ast.FunctionDef.
-            has_father_class (bool): whether class has father class that is not nn.Cell
-
-        Raises:
-            RuntimeError: Not support multi-targets in assign.
-            RuntimeError: Only support target.value in [ast.Name] in assign node.
-        """
-        new_bodies = []
-        for body in bodies:
-            if self._is_super_expr(body):
-                if has_father_class:
-                    body.value.args = [ast.Name(id='obj', ctx=ast.Load())]
-                    body.value.keywords = []
-                new_bodies.append(body)
-                continue
-        ast_copy_attr = ast.parse(
-            "for key, value in obj.__dict__.items(): setattr(self, key, value)").body[0]
-        new_bodies.append(ast_copy_attr)
-        return new_bodies
-
-    def _handle_father_class(self, stree, node: ast.ClassDef) -> bool:
-        """Handle father class."""
-        has_father_class = False
-        for base in node.bases:
-            parser: Parser = ParserRegister.instance().get_parser(type(base))
-            father_class = parser.process(stree, base)
-            if "Cell" not in father_class:
-                for k, m in sys.modules.items():
-                    if k in ("_ast", "ast"):
-                        continue
-                    if hasattr(m, father_class):
-                        cls = getattr(m, father_class)
-                        if not inspect.isclass(cls):
-                            continue
-                        source_code = inspect.getsource(cls)
-                        father_class_ast: ast.Module = ast.parse(source_code)
-                        self._father_class_process_init_func_ast(stree, father_class_ast)
-                        stree._father_class_ast.append(father_class_ast) # pylint: disable=protected-access
-                        has_father_class = True
-                        break
-        return has_father_class
-
-    def _father_class_process_init_func_ast(self, stree: SymbolTree, father_class_ast: ast.Module):
-        father_class_stree: SymbolTree = SymbolTree(stree.get_origin_network(), father_class_ast)
-        for ast_body in father_class_ast.body:
-            if isinstance(ast_body, ast.ClassDef):
-                has_father_class = self._handle_father_class(stree, ast_body)
-                if self._need_add_init_func(father_class_stree, ast_body):
-                    self._add_init_func(ast_body)
-                for body in ast_body.body:
-                    if isinstance(body, ast.FunctionDef) and body.name == "__init__":
-                        self._process_init_func_ast(body, has_father_class)
-
+        # Copy function class variables into new network
+        ClassDefParser._process_class_variables(stree, function_defs)
 
 g_classdef_parser = reg_parser(ClassDefParser())

mindspore/rewrite/parsers/constant_parser.py

@@ -15,11 +15,11 @@
 """Parse ast.Assign in construct function to node of SymbolTree."""
 import ast
 
-from mindspore.rewrite.parser import Parser
+from mindspore.rewrite.parsers.parser import Parser
 from mindspore.rewrite.symbol_tree import SymbolTree
-from mindspore.rewrite.parser_register import reg_parser
+from mindspore.rewrite.parsers.parser_register import reg_parser
 from ..common import error_str
-
+from ..node.node_manager import NodeManager
 
 class NameParser(Parser):
     """Parse ast.Name in construct function to node of SymbolTree."""
@@ -28,13 +28,14 @@ class NameParser(Parser):
         """Parse target type."""
         return ast.Name
 
-    def process(self, stree: SymbolTree, node: ast.Name):
+    def process(self, stree: SymbolTree, node: ast.Name, node_manager: NodeManager):
         """
         Parse ast.Name node.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([ast.Name]): An ast.Name node.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Raises:
             TypeError: Name parser only supports parsing ast.Name type nodes.
@@ -52,13 +53,14 @@ class NumParser(Parser):
         """Parse target type."""
         return ast.Num
 
-    def process(self, stree: SymbolTree, node: ast.Num):
+    def process(self, stree: SymbolTree, node: ast.Num, node_manager: NodeManager):
         """
         Parse ast.Num node.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([ast.Num]): An ast.Num node.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Raises:
             TypeError: Num parser only supports parsing ast.Num type nodes.
@@ -76,13 +78,14 @@ class StrParser(Parser):
         """Parse target type."""
         return ast.Str
 
-    def process(self, stree: SymbolTree, node: ast.Str):
+    def process(self, stree: SymbolTree, node: ast.Str, node_manager: NodeManager):
         """
         Parse ast.Str node.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([ast.Str]): An ast.Str node.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Returns:
             The value of node.

mindspore/rewrite/parsers/container_parser.py

@@ -15,12 +15,12 @@
 """Parse Container in construct function to node of SymbolTree."""
 import ast
 
-from mindspore.rewrite.parser import Parser
+from mindspore.rewrite.parsers.parser import Parser
 from mindspore.rewrite.symbol_tree import SymbolTree
-from mindspore.rewrite.parser_register import ParserRegister
+from mindspore.rewrite.parsers.parser_register import ParserRegister, reg_parser
 
-from mindspore.rewrite.parser_register import reg_parser
 from ..common import error_str
+from ..node.node_manager import NodeManager
 
 
 class ListParser(Parser):
@@ -30,13 +30,14 @@ class ListParser(Parser):
         """Parse target type."""
         return list
 
-    def process(self, stree: SymbolTree, node: list):
+    def process(self, stree: SymbolTree, node: list, node_manager: NodeManager):
         """
         Parse list.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([list]): An list of node.
+            father_node_managernode (NodeManager): NodeManager those asts belong to.
 
         Returns:
             A list of value.
@@ -50,7 +51,7 @@ class ListParser(Parser):
         result = []
         for n in node:
             parser = ParserRegister.instance().get_parser(type(n))
-            value = parser.process(stree, n)
+            value = parser.process(stree, n, node_manager)
             result.append(value)
         return result
 
@@ -62,13 +63,14 @@ class TupleParser(Parser):
         """Parse target type."""
         return tuple
 
-    def process(self, stree: SymbolTree, node: tuple):
+    def process(self, stree: SymbolTree, node: tuple, node_manager: NodeManager):
         """
         Parse tuple.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([tuple]): An tuple of node.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Returns:
             A tuple of value.
@@ -79,7 +81,7 @@ class TupleParser(Parser):
         result = []
         for n in node:
             parser = ParserRegister.instance().get_parser(type(n))
-            value = parser.process(stree, n)
+            value = parser.process(stree, n, node_manager)
             result.append(value)
         return tuple(result)
 

mindspore/rewrite/parsers/for_parser.py

@@ -13,17 +13,23 @@
 # limitations under the License.
 # ============================================================================
 """ Parse ast.For node """
+import sys
 import ast
-import astunparse
 
 from mindspore.rewrite.api.scoped_value import ScopedValue, ValueType
 from mindspore.rewrite.ast_helpers.ast_modifier import AstModifier
 from mindspore import log as logger
 from mindspore import nn
 from ..symbol_tree import SymbolTree
-from ..parser import Parser
-from ..parser_register import reg_parser
+from .parser import Parser
+from .parser_register import reg_parser
 from ..common.event import Event
+from ..node.node_manager import NodeManager
+
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
 
 EVAL_WHITE_LIST = ("self.", "range(", "zip(", "enumerate(", "reversed(")
 
@@ -34,20 +40,20 @@ class ForParser(Parser):
     @staticmethod
     def modify_init_ast(stree, i, obj, iter_var_name):
         """Modify the ast node in init function."""
-        target = f"{iter_var_name.strip()}_{str(i)}"
+        target = f"{iter_var_name.strip()}{str(i)}"
         setattr(stree.get_origin_network(), target, obj)
         stree.get_origin_network().insert_child_to_cell(target, obj)
         AstModifier.insert_assign_to_function(stree.get_init_func_ast(),
                                               targets=[ScopedValue(ValueType.NamingValue, "self", target)],
                                               expr=ScopedValue(ValueType.NamingValue, "", "getattr"),
                                               args=[ScopedValue(ValueType.NamingValue, "", "obj"),
-                                                    ScopedValue(ValueType.StringValue, "", target)])
+                                                    ScopedValue(ValueType.ConstantValue, "", target)])
 
     @staticmethod
     def modify_construct_ast(stree, ast_node, old_name, new_name):
         """Modify the ast node in construct function."""
         node_str: str = astunparse.unparse(ast_node)
-        node_str = node_str.replace(old_name, new_name)
+        node_str = node_str.replace(old_name+'(', new_name+'(')
         module_node = ast.parse(node_str)
         new_node = module_node.body[0]
         return new_node
@@ -55,10 +61,15 @@ class ForParser(Parser):
     def target(self):
         return ast.For
 
-    def process(self, stree: SymbolTree, node: ast.For):
+    def process(self, stree: SymbolTree, node: ast.For, node_manager: NodeManager):
         """ Process ast.For node """
         if isinstance(node.target, ast.Name):
             targets = node.target.id
+        if isinstance(node.iter, ast.Str) or (isinstance(node.iter, ast.Constant) and
+                                              isinstance(node.iter.val, str)):
+            # Ast.For which has iter with type of str is converted to python node to avoid instruction injection
+            stree.try_append_python_node(node, node)
+            return
         iter_code = astunparse.unparse(node.iter)
         if not iter_code.startswith(EVAL_WHITE_LIST):
             logger.warning(
@@ -72,26 +83,36 @@ class ForParser(Parser):
             _info = f"For MindSpore Rewrtie, when eval '{iter_code}' by using JIT Fallback feature, " \
                          f"an error occurred: {str(e)}"
             logger.warning(_info)
-            stree.try_append_python_node(node, node)
+            stree.try_append_python_node(node, node, node_manager)
             return
 
         iter_var_name = iter_code.split(".")[-1]
-        index = stree.get_ast_root().body.index(node) + 1
-        if isinstance(iter_obj, list):
+        ast_functiondef = node_manager.get_ast_functiondef()
+        if not ast_functiondef:
+            logger.info(f"ast_functiondef is None in node_manager {node_manager.get_manager_name()} "
+                        "when parsing 'for' statement.")
+            stree.try_append_python_node(node, node, node_manager)
+            return
+        index = ast_functiondef.body.index(node) + 1
+        if isinstance(iter_obj, (list, nn.CellList)):
             for obj in iter_obj:
                 if not isinstance(obj, nn.Cell):
-                    stree.try_append_python_node(node, node)
+                    stree.try_append_python_node(node, node, node_manager)
                     return
             for i, obj in enumerate(iter_obj):
                 ForParser.modify_init_ast(stree, i, obj, iter_var_name)
                 for body in node.body:
-                    new_func_name = f"self.{iter_var_name.strip()}_{str(i)}".strip()
+                    new_func_name = f"self.{iter_var_name.strip()}{str(i)}".strip()
                     new_node = ForParser.modify_construct_ast(stree, body, targets, new_func_name)
-                    stree.get_ast_root().body.insert(index, new_node)
+                    ast_functiondef.body.insert(index, new_node)
                     index += 1
+            # Expand "for" statement and replace the body with Pass
+            for body in node.body[:]:
+                node.body.remove(body)
+            node.body.append(ast.Pass())
+
             if stree.get_ori_cls_name() == "SequentialCell":
                 stree.on_change(Event.CodeChangeEvent)
-            stree.get_ast_root().body.remove(node)
             return
         if isinstance(iter_obj, range):
             logger.warning("For MindSpore Rewrite, range not support.")
@@ -101,7 +122,7 @@ class ForParser(Parser):
             logger.warning("For MindSpore Rewrite, enumerate not support.")
         else:
             logger.warning(f"For MindSpore Rewrite, not supported type: {type(iter_obj).__name__}")
-        stree.try_append_python_node(node, node)
+        stree.try_append_python_node(node, node, node_manager)
         return
 
 g_for_parser = reg_parser(ForParser())