mindspore 2.1.0__cp39-none-any.whl → 2.2.11__cp39-none-any.whl

mindspore/rewrite/ast_helpers/ast_modifier.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """Ast utils for create or update ast node."""
-from typing import Optional
+from typing import Optional, List
 import ast
 
 from ..api.scoped_value import ScopedValue, ValueType
@@ -34,20 +34,15 @@ class AstModifier(ast.NodeTransformer):
         Returns:
             A bool if to_erase-node been found and been erased.
         """
-        for body in ast_func.body:
+        return AstModifier.erase_ast_from_bodies(ast_func.body, to_erase)
+
+    @staticmethod
+    def erase_ast_from_bodies(ast_bodies: List[ast.AST], to_erase: ast.AST) -> bool:
+        """Erase ast node from ast bodies."""
+        for body in ast_bodies:
             if id(body) == id(to_erase):
-                ast_func.body.remove(body)
+                ast_bodies.remove(body)
                 return True
-            # hardcode for ast.If
-            if isinstance(body, ast.If):
-                for if_body in body.body:
-                    if id(if_body) == id(to_erase):
-                        body.body.remove(if_body)
-                        return True
-                for else_body in body.orelse:
-                    if id(else_body) == id(to_erase):
-                        body.orelse.remove(else_body)
-                        return True
         return False
 
     @staticmethod
@@ -147,13 +142,6 @@ class AstModifier(ast.NodeTransformer):
             RuntimeError: If 'index_ast' is not contained in 'ast_func'.
         """
         assign = AstModifier.create_call_assign(targets, expr, args, kwargs)
-        arguments: ast.arguments = ast_func.args
-        if arguments.args:
-            for arg in arguments.args:
-                if id(arg) == id(index_ast):
-                    ast_func.body.insert(0, assign)
-                    ast.fix_missing_locations(ast_func)
-                    return assign
         return AstModifier.insert_assign_ast_to_function(ast_func, assign, index_ast, insert_before)
 
     @staticmethod
@@ -177,49 +165,63 @@ class AstModifier(ast.NodeTransformer):
         Raises:
             RuntimeError: If 'index_ast' is not contained in 'ast_func'.
         """
-        if index_ast is None:
-            ast_func.body.append(ast_assign)
-            ast.fix_missing_locations(ast_func)
-            return ast_assign
+        # Insert ast at the frontmost position of function body when index_ast is an argument of function
         arguments: ast.arguments = ast_func.args
-        if arguments.args:
+        if index_ast and arguments.args:
             for arg in arguments.args:
                 if id(arg) == id(index_ast):
                     ast_func.body.insert(0, ast_assign)
                     ast.fix_missing_locations(ast_func)
                     return ast_assign
-        for index in range(0, len(ast_func.body)):
-            body = ast_func.body[index]
+        # Insert ast at position specified by index_ast in function body
+        ast_assign = AstModifier.insert_assign_ast_to_bodies(ast_func.body, ast_assign, index_ast, insert_before)
+        ast.fix_missing_locations(ast_assign)
+        return ast_assign
+
+    @staticmethod
+    def insert_assign_ast_to_bodies(ast_bodies: List[ast.AST], ast_assign: ast.Assign,
+                                    index_ast: Optional[ast.AST] = None, insert_before=True) -> ast.AST:
+        """Insert ast at position specified by index_ast of ast_bodies"""
+        # Append ast_assign to ast_bodies when index_ast is None
+        if index_ast is None:
+            ast_bodies.append(ast_assign)
+            return ast_assign
+        # Append ast_assign to ast_bodies
+        for index, body in enumerate(ast_bodies):
             if id(body) == id(index_ast):
-                if insert_before:
-                    ast_func.body.insert(index, ast_assign)
-                else:
-                    ast_func.body.insert(index + 1, ast_assign)
-                ast.fix_missing_locations(ast_func)
-                return ast_assign
-            # hardcode for ast.If
-            if isinstance(body, ast.If):
-                for if_index in range(0, len(body.body)):
-                    if_body = body.body[if_index]
-                    if id(if_body) != id(index_ast):
-                        continue
-                    if insert_before:
-                        body.body.insert(if_index, ast_assign)
-                    else:
-                        body.body.insert(if_index + 1, ast_assign)
-                    ast.fix_missing_locations(body)
-                    return ast_assign
-                for if_index in range(0, len(body.orelse)):
-                    else_body = body.orelse[if_index]
-                    if id(else_body) != id(index_ast):
-                        continue
-                    if insert_before:
-                        body.orelse.insert(if_index, ast_assign)
-                    else:
-                        body.orelse.insert(if_index + 1, ast_assign)
-                    ast.fix_missing_locations(body)
-                    return ast_assign
-        raise RuntimeError("insert position is not contained in ast_func")
+                if not insert_before:
+                    index += 1
+                ast_bodies.insert(index, ast_assign)
+                ast.fix_missing_locations(body)
+                break
+        else:
+            raise ValueError("insert position is not contained in ast_bodies")
+        return ast_assign
+
+    @staticmethod
+    def append_arg_to_function(ast_func: ast.FunctionDef, ast_arg: ast.arg) -> ast.AST:
+        """
+        Append an ast.arg to an ast.FunctionDef (e.g. self.construct).
+
+        Args:
+            ast_func (ast.FunctionDef): An instance of ast.FunctionDef which is "construct" function of network.
+            ast_arg (ast.arg): An instance of ast.arg to be inserted in.
+
+        Returns:
+            An instance of ast.arg which has been appended to 'ast_func'.
+
+        Raises:
+            RuntimeError: If 'ast_arg' is not an instance of ast_arg.
+        """
+        if not isinstance(ast_arg, ast.arg):
+            raise RuntimeError("ast_arg should be an instance of ast.arg.")
+        arguments: ast.arguments = ast_func.args
+        args: [ast.arg] = arguments.args
+        args.append(ast_arg)
+        defaults = arguments.defaults
+        arg_default = ast.Constant(value=None, kind=None)
+        defaults.append(arg_default)
+        return ast_arg
 
     @staticmethod
     def append_global_vars_expr_to_init(init_func: ast.FunctionDef, targets: [ScopedValue],
@@ -241,7 +243,7 @@ class AstModifier(ast.NodeTransformer):
             An instance of ast.Assign which has been appended to 'init_func'.
         """
         return AstModifier.insert_assign_to_function(init_func, targets=targets,
-                                                     expr=ScopedValue(ValueType.NamingValue, "", "setattr"),
+                                                     expr=ScopedValue(ValueType.NamingValue, "", "getattr"),
                                                      args=[ScopedValue(ValueType.NamingValue, "obj"),
                                                            ScopedValue.create_variable_value(field)])
 
@@ -265,23 +267,30 @@ class AstModifier(ast.NodeTransformer):
             RuntimeError: If 'targets' is None.
             RuntimeError: If value_type of element of 'targets' is not ValueType.NamingValue.
 
-            RuntimeError: If length of 'targets' is not 1. Multi-targets will be support in the future.
         """
-        if targets is None or len(targets) != 1:
-            raise RuntimeError("Only support one target in insert_cell_to_init now")
-        if targets[0].type != ValueType.NamingValue:
-            raise RuntimeError("Target must be a right-value, got: ", targets[0])
-        if targets[0].scope:
-            ast_target = ast.Attribute(ast.Name(targets[0].scope, ast.Load()), targets[0].value, ast.Store())
-        else:
-            ast_target = ast.Name(targets[0].value, ast.Store())
+        if targets is None:
+            raise RuntimeError("'Targets should not be None.")
+        targets_list = []
+        for target in targets:
+            if target.type != ValueType.NamingValue:
+                raise RuntimeError("Target must be a right-value, got: ", target)
+            if target.scope:
+                ast_target = ast.Attribute(ast.Name(target.scope, ast.Load()), target.value, ast.Store())
+            else:
+                ast_target = ast.Name(target.value, ast.Store())
+            targets_list.append(ast_target)
         call = AstModifier.create_call(expr, args, kwargs)
-        result = ast.Assign(targets=[ast_target], value=call)
+
+        if len(targets) == 1:
+            result = ast.Assign(targets=[targets_list[0]], value=call)
+        elif len(targets) > 1:
+            ast_targets = ast.Tuple(elts=targets_list, ctx=ast.Store())
+            result = ast.Assign(targets=[ast_targets], value=call)
         ast.fix_missing_locations(result)
         return result
 
     @staticmethod
-    def _create_arg_by_single_value(value: ScopedValue):
+    def _create_arg_by_constant_value(value: ScopedValue):
         """
         Create an instance of ast.Constant.
 
@@ -290,17 +299,16 @@ class AstModifier(ast.NodeTransformer):
 
         Raises:
             RuntimeError: if scope of value is not empty.
-            TypeError: type of arg not in [ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue]
+            TypeError: type of arg is not ValueType.ConstantValue
 
         Returns:
             ast.Constant: An instance of ast.Constant
         """
-        if value.type in (ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue):
+        if value.type == ValueType.ConstantValue:
             if value.scope:
                 raise RuntimeError("For arg the scope should be empty")
             return ast.Constant(value=value.value, kind=None)
-        raise TypeError("Type of arg only support [ValueType.IntValue, ValueType.FloatValue,"
-                        f" ValueType.StringValue], but got {type(value)}")
+        raise TypeError("Type of arg only support ValueType.ConstantValue, but got {type(value)}")
 
     @staticmethod
     def _create_list_or_tuple(value: ScopedValue):
@@ -315,7 +323,7 @@ class AstModifier(ast.NodeTransformer):
         """
         elts = []
         for v in value.value:
-            elts.append(AstModifier._create_arg_by_single_value(v))
+            elts.append(AstModifier._create_arg_by_constant_value(v))
         if isinstance(value, list):
             return ast.List(elts=elts)
         return ast.Tuple(elts=elts)
@@ -331,22 +339,20 @@ class AstModifier(ast.NodeTransformer):
 
         Raises:
             RuntimeError:  if scope of value is not empty.
-            TypeError: type of arg not in [ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue,
-            ValueType.ListValue, ValueType.TupleValue]
+            TypeError: type of arg is not ValueType.ConstantValue
 
         Returns:
             ast.keyword: a instance of ast.keyword.
         """
         if value.scope:
             raise RuntimeError("value.scope should be empty")
-        if value.type in (ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue):
+        if value.type == ValueType.ConstantValue:
             v = ast.Constant(value=value.value, kind=None)
         elif value.type in (ValueType.ListValue, ValueType.TupleValue):
             v = AstModifier._create_list_or_tuple(value)
         else:
-            raise TypeError("Type of keyword value only support [ValueType.IntValue, ValueType.FloatValue,"
-                            "ValueType.StringValue, ValueType.ListValue, ValueType.TupleValue],"
-                            f"but got {type(value)}")
+            raise TypeError("Type of keyword value only support [ValueType.ConstantValue, ValueType.ListValue, "
+                            f"ValueType.TupleValue], but got {type(value)}")
         return ast.keyword(arg=arg, value=v)
 
     @staticmethod
@@ -371,14 +377,14 @@ class AstModifier(ast.NodeTransformer):
         for arg in args:
             if not isinstance(arg, ScopedValue):
                 raise TypeError("arg should be ScopedValue, got: ", type(arg))
-            if arg.type in (ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue):
-                results.append(AstModifier._create_arg_by_single_value(arg))
+            if arg.type == ValueType.ConstantValue:
+                results.append(AstModifier._create_arg_by_constant_value(arg))
             elif arg.type == ValueType.NamingValue:
                 if arg.scope:
                     results.append(ast.Attribute(ast.Name(arg.scope, ast.Load()), arg.value, ast.Store()))
                 else:
                     results.append(ast.Name(arg.value, ast.Store()))
-            elif arg.type == ValueType.ListValue or arg.type == ValueType.TupleValue:
+            elif arg.type in (ValueType.ListValue, ValueType.TupleValue):
                 results.append(AstModifier._create_list_or_tuple(arg))
             else:
                 raise RuntimeError("Please handle custom-object first")
@@ -406,8 +412,7 @@ class AstModifier(ast.NodeTransformer):
         for arg, value in kwargs.items():
             if not isinstance(value, ScopedValue):
                 raise TypeError("value should be ScopedValue, got: ", type(value))
-            if value.type in (ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue,
-                              ValueType.ListValue, ValueType.TupleValue):
+            if value.type in (ValueType.ConstantValue, ValueType.ListValue, ValueType.TupleValue):
                 results.append(AstModifier._create_keyword(arg, value))
             elif value.type == ValueType.NamingValue:
                 if value.scope:
@@ -466,7 +471,7 @@ class AstModifier(ast.NodeTransformer):
         Raises:
             TypeError: Input src_argument is not a ScopedValue
             RuntimeError: If 'dst_ast' is an instance of ast.Constant but type of 'src_argument' is not
-                ValueType.IntValue, ValueType.FloatValue or ValueType.StringValue.
+                ValueType.ConstantValue.
             RuntimeError: If 'dst_ast' is an instance of ast.Name or ast.Attribute but type of 'src_argument' is not
                 ValueType.NamingValue.
             RuntimeError: When 'dst_ast' is an instance of ast.Name, scope of 'src_argument' is not empty.
@@ -480,27 +485,14 @@ class AstModifier(ast.NodeTransformer):
         """
         if not isinstance(src_argument, ScopedValue):
             raise TypeError("src_argument should be ScopedValue, got: ", type(src_argument))
-        if isinstance(dst_ast, ast.Constant):
-            if src_argument.type not in [ValueType.IntValue, ValueType.FloatValue, ValueType.StringValue]:
-                raise RuntimeError("src_argument should be a IntValue, FloatValue or StringValue, got:",
-                                   str(src_argument.type))
-            dst_ast.value = src_argument.value
-            return
-        if isinstance(dst_ast, ast.Num):
-            if src_argument.type not in [ValueType.IntValue, ValueType.FloatValue]:
-                raise RuntimeError("src_argument should be a IntValue or FloatValue, but got:",
-                                   str(src_argument.type))
-            dst_ast.n = src_argument.value
-            return
-        if isinstance(dst_ast, ast.Str):
-            if src_argument.type not in [ValueType.StringValue]:
-                raise RuntimeError("src_argument should be a StringValue, but got:",
-                                   str(src_argument.type))
-            dst_ast.s = src_argument.value
+        if isinstance(dst_ast, (ast.Constant, ast.Num, ast.Str)):
+            AstModifier.update_arg_value_constant(src_argument, dst_ast)
             return
         if isinstance(dst_ast, ast.Name):
-            if src_argument.type not in [ValueType.NamingValue, ValueType.StringValue]:
-                raise RuntimeError("src_argument.type should be ValueType.NamingValue or ValueType.StringValue.")
+            if src_argument.type not in [ValueType.NamingValue, ValueType.ConstantValue]\
+                or not isinstance(src_argument.value, str):
+                raise RuntimeError("src_argument.type should be ValueType.NamingValue or ValueType.ConstantValue, "
+                                   "but got:", type(src_argument.value).__name__)
             if src_argument.scope:
                 raise RuntimeError("src_argument.scope should be empty")
             dst_ast.id = src_argument.value
@@ -523,3 +515,23 @@ class AstModifier(ast.NodeTransformer):
                 AstModifier.update_arg_value(src_argument.value[elt_index], elt)
             return
         raise RuntimeError("keyword value type is not supported", type(dst_ast))
+
+    @staticmethod
+    def update_arg_value_constant(src_argument: ScopedValue, dst_ast: ast.AST):
+        """Update 'arg_value' of type constant by 'input_argument'"""
+        if src_argument.type != ValueType.ConstantValue:
+            raise RuntimeError("src_argument should be a ConstantValue, got:", str(src_argument.type))
+        if isinstance(dst_ast, ast.Constant):
+            dst_ast.value = src_argument.value
+            return
+        if isinstance(dst_ast, ast.Num):
+            if not isinstance(src_argument.value, (int, float)):
+                raise RuntimeError("Type of src_argument should be int or float, but got:",
+                                   type(src_argument.value).__name__)
+            dst_ast.n = src_argument.value
+            return
+        if isinstance(dst_ast, ast.Str):
+            if not isinstance(src_argument.value, str):
+                raise RuntimeError("Type of src_argument should be str, but got:", type(src_argument.value).__name__)
+            dst_ast.s = src_argument.value
+            return

mindspore/rewrite/ast_transformers/flatten_recursive_stmt.py

@@ -14,14 +14,20 @@
 # ============================================================================
 """Ast optimizer for flatten recursive call."""
 
-from typing import Any, Tuple
+import sys
+from typing import Any, Tuple, List
+import keyword
 import ast
-from ast import FunctionDef
-import astunparse
 
 from mindspore import log as logger
 from ..common import error_str
 
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
+
+FLATTEN_BLACK_LIST = ["set_vertex_attr",]
 
 class FlattenRecursiveStmt(ast.NodeTransformer):
     """Ast optimizer for flatten recursive call."""
@@ -40,17 +46,35 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
             ast.BoolOp: ["values"],
             ast.UnaryOp: ["operand"],
             ast.Compare: ["left", "comparators"],
+            ast.If: ["test"]
         }
+        self._transform_functions = []
+        self._transform_if = False
+        self._symbol_tree = None # Used to get unique name
 
     @staticmethod
-    def _generate_target_name(node: ast.AST, target_names):
+    def _check_flatten_black_list(node: ast.AST):
+        """Check whether node in flatten black list"""
+        func_name = ""
+        # Get func name of node
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name):
+                func_name = node.func.id
+            elif isinstance(node.func, ast.Attribute):
+                func_name = node.func.attr
+        # Check func name of node
+        if func_name and func_name in FLATTEN_BLACK_LIST:
+            return True
+        return False
+
+    def _generate_target_name(self, node: ast.AST, target_names):
         """Generate unique target name."""
         if isinstance(node, ast.Call):
             func = node.func
             if isinstance(func, ast.Name):
-                target_name = func.id
+                target_name = func.id + "_var"
             elif isinstance(func, ast.Attribute):
-                target_name = func.attr
+                target_name = func.attr + "_var"
             else:
                 logger.info("unhandled type of func of ast.Call while generating new target name: %s ", type(func))
                 target_name = "function"
@@ -67,30 +91,33 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
         else:
             logger.info("unhandled type of node while generating new target name: %s ", type(node))
             target_name = type(node).__name__.lower() + "_var"
+        # avoid python keyword
+        if keyword.iskeyword(target_name):
+            target_name = target_name + "_var"
         suffix = 0
         result = target_name
         while result in target_names:
             suffix += 1
             result = f"{target_name}_{suffix}"
+        if self._symbol_tree:
+            result = self._symbol_tree.unique_name(result)
         target_names.append(result)
         return result
 
-    @staticmethod
-    def _create_new_assign_node(node: ast.AST, target_names) -> Tuple[str, ast.AST]:
+    def _create_new_assign_node(self, node: ast.AST, target_names) -> Tuple[str, ast.AST]:
         """Create new assign node to be inserted into ast.FunctionDef."""
         if isinstance(node, (ast.Name, ast.Constant, ast.Num, ast.Str, ast.NameConstant, ast.Bytes, ast.Ellipsis)):
             return "", node
-        new_target_name = FlattenRecursiveStmt._generate_target_name(node, target_names)
+        new_target_name = self._generate_target_name(node, target_names)
         return new_target_name, ast.Assign(targets=[ast.Name(id=new_target_name, ctx=ast.Store())], value=node)
 
-    @staticmethod
-    def _flatten_list(node_list, target_names):
+    def _flatten_list(self, node_list, target_names):
         """Flatten a list of node."""
         results = list()
         new_list = list()
         for node in node_list:
             if isinstance(node, ast.Call):
-                new_target, new_node = FlattenRecursiveStmt._create_new_assign_node(node, target_names)
+                new_target, new_node = self._create_new_assign_node(node, target_names)
                 results.append(new_node)
                 new_list.append(ast.Name(id=new_target, ctx=ast.Load()))
             else:
@@ -99,6 +126,8 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
 
     def _flatten_statement(self, node: ast.AST, target_names) -> [ast.AST]:
         """Flatten recursive statement according to different node type."""
+        if FlattenRecursiveStmt._check_flatten_black_list(node):
+            return []
         flatten_config = self._flatten_table.get(type(node))
         if flatten_config is None:
             return []
@@ -112,21 +141,21 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                     if isinstance(todo, ast.Starred):
                         new_list.append(todo)
                         continue
-                    new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(todo, target_names)
+                    new_target_name, new_node = self._create_new_assign_node(todo, target_names)
                     if id(new_node) == id(todo):
                         new_list.append(todo)
                     else:
                         new_list.append(ast.Name(id=new_target_name, ctx=ast.Load()))
                         results.append(new_node)
                     if isinstance(todo, (ast.Tuple, tuple)):
-                        _res, _new_list = FlattenRecursiveStmt._flatten_list(new_node.value.elts, [new_target_name])
+                        _res, _new_list = self._flatten_list(new_node.value.elts, [new_target_name])
                         new_node.value.elts = _new_list
                         results.extend(_res)
                 setattr(node, todo_name, new_list)
             elif isinstance(todos, dict):
                 new_dict = []
                 for key, value in todos:
-                    new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(value, target_names)
+                    new_target_name, new_node = self._create_new_assign_node(value, target_names)
                     if id(new_node) == id(value):
                         new_dict[key] = value
                     else:
@@ -134,16 +163,15 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                         results.append(new_node)
                 setattr(node, todo_name, new_dict)
             else:
-                new_target_name, new_node = FlattenRecursiveStmt._create_new_assign_node(todos, target_names)
+                new_target_name, new_node = self._create_new_assign_node(todos, target_names)
                 if id(new_node) != id(todos):
                     setattr(node, todo_name, ast.Name(id=new_target_name, ctx=ast.Load()))
                     results.append(new_node)
         return results
 
-    def _fill_in_original_target_names(self, target_names, node):
-        """Fill in original target names before getting unique names."""
-        for function_index in range(len(node.body)):
-            child = node.body[function_index]
+    def _save_target_names(self, target_names, ast_body: List[ast.AST]):
+        """Saving target names in ast_body before getting unique names."""
+        for child in ast_body:
             if isinstance(child, (ast.Assign, ast.Expr)):
                 child_value = child.value
             else:
@@ -155,7 +183,7 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                 continue
             targets = child.targets
             for target in targets:
-                if not isinstance(target, (ast.Name, ast.Tuple)):
+                if not isinstance(target, (ast.Name, ast.Tuple, ast.List)):
                     raise RuntimeError(
                         error_str(f"currently only support ast.Name targets, but got ast type "
                                   f"'{type(target).__name__}'", child_node=target, father_node=child))
@@ -163,7 +191,7 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                     target_name = target.id
                     if target_name not in target_names:
                         target_names.append(target_name)
-                elif isinstance(target, ast.Tuple):
+                elif isinstance(target, (ast.Tuple, ast.List)):
                     for elt in target.elts:
                         if not isinstance(elt, ast.Name):
                             raise RuntimeError(
@@ -174,47 +202,66 @@ class FlattenRecursiveStmt(ast.NodeTransformer):
                         if target_name not in target_names:
                             target_names.append(target_name)
 
-    def visit_FunctionDef(self, node: FunctionDef) -> Any:
-        """Traverse construct node and flatten recursive nodes."""
-        if node.name != "construct":
-            return node
-
+    def _visit_ast_bodies(self, ast_body: List[ast.AST]):
+        """Traverse nodes in ast_body and flatten nodes recursive."""
         target_names = []
-        self._fill_in_original_target_names(target_names, node)
-        index = len(node.body) - 1
+        self._save_target_names(target_names, ast_body)
+        index = len(ast_body) - 1
         while index >= 0:
-            child = node.body[index]
+            child = ast_body[index]
             if isinstance(child, ast.Assign):
                 stmt = child.value
             elif isinstance(child, ast.If):
                 if isinstance(child.body[0], ast.Return) and not isinstance(child.test, ast.UnaryOp):
-                    if isinstance(child.body[0].value, ast.Call):
-                        if_body = child.body
-                        if_func = if_body[0].value
-                        expr = "x = " + astunparse.unparse(if_func)
-                        if_body = ast.parse(expr)
-                        if_body = if_body.body+ast.parse("return x").body
-                        child.body = if_body
-                        stmt = child
-                    else:
-                        stmt = child
-                else:
-                    stmt = child
+                    if not isinstance(child.body[0].value, (ast.Name, ast.Constant)):
+                        return_val_ast = child.body[0].value
+                        return_name = self._generate_target_name(return_val_ast, target_names)
+                        new_assign_code = f"{return_name} = {astunparse.unparse(return_val_ast)}"
+                        new_assign_ast = ast.parse(new_assign_code).body[0]
+                        new_return_ast = ast.parse(f"return {return_name}").body[0]
+                        child.body = [new_assign_ast, new_return_ast]
+                stmt = child
             elif isinstance(child, ast.Expr):
                 stmt = child.value
             else:
                 stmt = child
             results = self._flatten_statement(stmt, target_names)
             if results:
-                results.reverse()
-                for result in results:
-                    node.body.insert(index, result)
+                for result in reversed(results):
+                    ast_body.insert(index, result)
                     index += 1
             index -= 1
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: # pylint: disable=invalid-name
+        """Traverse nodes in _transform_functions and flatten recursive nodes."""
+        if node.name not in self._transform_functions:
+            return node
+        self._visit_ast_bodies(node.body)
+        return node
+
+    def visit_If(self, node: ast.If) -> Any: # pylint: disable=invalid-name
+        """Traverse nodes in if node and flatten recursive nodes."""
+        if not self._transform_if:
+            return node
+        self._visit_ast_bodies(node.body)
+        if node.orelse:
+            self._visit_ast_bodies(node.orelse)
         return node
 
-    def transform(self, ast_root):
+    def transform(self, ast_root, transform_functions=None, stree=None):
         """Interface of FlattenRecursiveStmt."""
+        self._transform_functions = transform_functions if transform_functions else ["construct"]
+        self._transform_if = False
+        self._symbol_tree = stree
         ast_root = self.visit(ast_root)
         ast_root = ast.fix_missing_locations(ast_root)
         return ast_root
+
+    def transform_if(self, ast_if, stree=None):
+        """Interface of FlattenRecursiveStmt."""
+        self._transform_functions = []
+        self._transform_if = True
+        self._symbol_tree = stree
+        ast_if = self.visit(ast_if)
+        ast_if = ast.fix_missing_locations(ast_if)
+        return ast_if

mindspore/rewrite/common/rewrite_elog.py

@@ -14,8 +14,12 @@
 # ============================================================================
 """Error Log for Rewrite."""
 
+import sys
 import ast
-import astunparse
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
 
 
 def error_str(reason: str, child_node: ast.expr = None, father_node: ast.expr = None) -> str: