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

mindspore/rewrite/parsers/function_def_parser.py

@@ -15,24 +15,25 @@
 """Parse bodies of ast.FunctionDef which is construct function to nodes of SymbolTree."""
 import ast
 from mindspore import log as logger
-from ..parser_register import ParserRegister, reg_parser
-from ..parser import Parser
+from .parser_register import ParserRegister, reg_parser
+from .parser import Parser
 from ..symbol_tree import SymbolTree
 from ..api.node_type import NodeType
+from ..node.node_manager import NodeManager
 
 
 class FunctionDefParser(Parser):
-    """Parse bodies of ast.FunctionDef which is construct function to nodes of SymbolTree."""
+    """Parse bodies of ast.FunctionDef in SymbolTree."""
 
     def target(self):
         """Parse target type"""
         return ast.FunctionDef
 
-    def remove_dead_code(self, stree: SymbolTree):
+    def remove_dead_code(self, node_manager: NodeManager):
         """Remove dead codes"""
         # Find out return node position
         return_idx = -1
-        for idx, node in enumerate(stree.nodes()):
+        for idx, node in enumerate(node_manager.nodes()):
             if node.get_node_type() == NodeType.Output:
                 return_idx = idx
                 break
@@ -40,29 +41,36 @@ class FunctionDefParser(Parser):
             return
         # Remove nodes after return node.
         # Reverse traversal to ensure that nodes are orphaned and can be deleted.
-        for idx, node in reversed(list(enumerate(stree.nodes()))):
+        for idx, node in reversed(list(enumerate(node_manager.nodes()))):
             if idx <= return_idx:
                 break
             logger.info(f"Remove dead code node:{node.get_name()}")
-            stree.erase_node(node)
+            node_manager.erase_node(node)
 
-    def process(self, stree: SymbolTree, node: ast.FunctionDef):
-        """Parse bodies of ast.FunctionDef which is construct function to nodes of SymbolTree."""
-        stree.set_ast_root(node)
+    def process(self, stree: SymbolTree, ast_node: ast.FunctionDef, node_manager: NodeManager):
+        """
+        Parse bodies of ast.FunctionDef in SymbolTree.
+
+        Args:
+            stree (SymbolTree): symbol tree under parsing.
+            ast_node (ast.FunctionDef): Ast FunctionDef node in construct.
+            node_manager (NodeManager): NodeManager those asts belong to.
+        """
         # parse args as inputs of stree
-        arguments: ast.arguments = node.args
+        arguments: ast.arguments = ast_node.args
         parser: Parser = ParserRegister.instance().get_parser(ast.arguments)
-        parser.process(stree, arguments)
+        parser.process(stree, arguments, node_manager)
 
         # parse body as node of stree
-        for body in node.body:
+        for body in ast_node.body:
             # avoid add dead code, so we need to break if return is added.
             parser: Parser = ParserRegister.instance().get_parser(type(body))
             if parser is None:
-                stree.append_python_node(node, body)
+                stree.append_python_node(ast_node, body, node_manager)
             else:
-                parser.process(stree, body)
-        self.remove_dead_code(stree)
+                parser.process(stree, body, node_manager)
+
+        self.remove_dead_code(node_manager)
 
 
 g_functiondef_parser = reg_parser(FunctionDefParser())

mindspore/rewrite/parsers/if_parser.py

@@ -15,11 +15,12 @@
 """Parse ast.If in construct function to node of SymbolTree."""
 
 import ast
-import astunparse
 
 from ..symbol_tree import SymbolTree
-from ..parser import Parser
-from ..parser_register import reg_parser
+from .parser import Parser
+from .parser_register import ParserRegister, reg_parser
+from ..node import NodeManager, ControlFlow
+from ..ast_transformers.flatten_recursive_stmt import FlattenRecursiveStmt
 
 
 class IfParser(Parser):
@@ -29,35 +30,38 @@ class IfParser(Parser):
         """Parse target type"""
         return ast.If
 
-    def process(self, stree: SymbolTree, node: ast.If):
+    def process(self, stree: SymbolTree, node: ast.If, node_manager: NodeManager):
         """
-        Parse ast.If and create a node in symbol tree.
+        Parse ast.If and create nodes into symbol tree.
 
         Args:
             stree ([SymbolTree]): Symbol Tree under parsing.
             node ([ast.If]): An ast.If node.
+            node_manager (NodeManager): NodeManager those asts belong to.
 
         Raises:
             NotImplementedError: If test of ast.If can not be eval.
         """
+        # expand codes in ast.if
+        ast_if = FlattenRecursiveStmt().transform_if(node, stree)
+        # parse ast codes of if branch into ControlFlow Node
+        if_node = ControlFlow("if_node", ast_if.body, stree)
+        for body in ast_if.body:
+            parser: Parser = ParserRegister.instance().get_parser(type(body))
+            if parser is None:
+                stree.append_python_node(ast_if, body, node_manager=if_node)
+            else:
+                parser.process(stree, body, node_manager=if_node)
+        stree.append_origin_field(if_node, node_manager)
+        # parse ast codes of else branch into ControlFlow Node
+        if ast_if.orelse:
+            else_node = ControlFlow("else_node", ast_if.orelse, stree)
+            for body in ast_if.orelse:
+                parser: Parser = ParserRegister.instance().get_parser(type(body))
+                if parser is None:
+                    stree.append_python_node(ast_if, body, node_manager=else_node)
+                else:
+                    parser.process(stree, body, node_manager=else_node)
+            stree.append_origin_field(else_node, node_manager)
 
-        test_code = astunparse.unparse(node.test)
-        test_code = test_code.replace("self", "stree.get_origin_network()")
-        bodies = None
-        try:
-            test_value = eval(test_code)
-        except (NameError, TypeError):
-            stree.try_append_python_node(node, node)
-            return
-
-        bodies = node.body if test_value else node.orelse
-        index = stree.get_ast_root().body.index(node) + 1
-        info_node = ast.Name(id=f"# If node has been replaced by {bool(test_value)} branch.",
-                             lineno=0, col_offset=0, ctx=ast.Load)
-        exp_node = ast.Expr(value=info_node, lineno=0, col_offset=0, ctx=ast.Load)
-        stree.get_ast_root().body.insert(index-1, exp_node)
-        for body in bodies:
-            stree.get_ast_root().body.insert(index, body)
-            index += 1
-        stree.get_ast_root().body.remove(node)
 g_if_parser = reg_parser(IfParser())

mindspore/rewrite/parsers/module_parser.py

@@ -13,23 +13,32 @@
 # limitations under the License.
 # ============================================================================
 """Parse ast.Module to SymbolTrees."""
+import sys
 from typing import Any
 import os
 import ast
 import copy
 import inspect
-import astunparse
 
 from mindspore import log as logger
 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 AstFinder
 from ..common import error_str
+from ..node.node_manager import NodeManager
 
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
 
 class ModuleParser(Parser):
     """Parse ast.Module to SymbolTrees."""
+
+    # a denied_class_decorator_list represents the decorators should be banned, which is registered by user
+    denied_class_decorator_list = []
+
     @staticmethod
     def _find_class(ast_node: ast.Module) -> ast.ClassDef:
         """Find all ast.ClassDef in ast.Module, only support one ast.ClassDef in ast.Module now."""
@@ -45,18 +54,27 @@ class ModuleParser(Parser):
     def _get_import_node(ast_root):
         """Iterate over ast_root and return all ast.Import nodes or ast.ImportFrom nodes in ast_root."""
         import_nodes = []
+        try_nodes = []
+        imports_str = []
 
         class GetImportNode(ast.NodeVisitor):
             """Find all import nodes from input ast node."""
 
+            def visit_Try(self, node: ast.Try) -> Any:
+                if isinstance(node.body[0], (ast.Import, ast.ImportFrom)):
+                    try_nodes.append(copy.deepcopy(node))
+                return node
+
             def visit_Import(self, node: ast.Import) -> Any:
                 """Iterate over all nodes and save ast.Import nodes."""
                 import_nodes.append(copy.deepcopy(node))
+                imports_str.append(astunparse.unparse(node))
                 return node
 
             def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:
                 """Iterate over all nodes and save ast.ImportFrom nodes."""
                 import_nodes.append(copy.deepcopy(node))
+                imports_str.append(astunparse.unparse(node))
                 return node
 
             def get_node(self, input_ast):
@@ -64,19 +82,145 @@ class ModuleParser(Parser):
                 self.generic_visit(input_ast)
                 return True
 
+        def _remove_duplicated_import_in_try(node: [ast.Import, ast.ImportFrom]):
+            import_str = astunparse.unparse(node)
+            if import_str in imports_str:
+                import_nodes.remove(import_nodes[imports_str.index(import_str)])
+
         get_node_handler = GetImportNode()
         get_node_handler.get_node(ast_root)
+        for Try in try_nodes:
+            for body in Try.body:
+                _remove_duplicated_import_in_try(body)
+            for handler in Try.handlers:
+                for body in handler.body:
+                    _remove_duplicated_import_in_try(body)
+        import_nodes.extend(try_nodes)
         return import_nodes
 
     @staticmethod
-    def _add_import_to_module(module: ast.Module, origin_net):
+    def save_file_path_to_sys(stree, level_num, file_path):
+        """
+        Save file path into stree._import_asts. `level_num` is used when level exist in ast.ImportFrom.
+
+        When level_num = 0(e.g. from xxx import yyy), current path will be saved.
+        When level_num = 1(e.g. from .xxx import yyy), current path will be saved.
+        When level_num = 2(e.g. from ..xxx import yyy), the path one level above the current path will be saved.
+        """
+        file_path = os.path.dirname(os.path.abspath(file_path))
+        if level_num > 1:
+            for _ in range(level_num - 1):
+                file_path = os.path.dirname(file_path)
+        sys_path_append_ast = ast.parse(f"sys.path.insert(0, r'{file_path}')").body[0]
+        stree.get_import_asts().append(ast.Import([ast.alias(name='sys', asname=None)]))
+        stree.get_import_asts().append(sys_path_append_ast)
+
+    @staticmethod
+    def _save_imports(stree):
         """Insert two groups of import nodes to ast.Module, common ones and those from class definition file."""
-        module.body.insert(0, ast.Import([ast.alias(name='mindspore', asname=None)]))
-        module.body.insert(1, ast.ImportFrom(module='mindspore', names=[ast.alias(name='nn', asname=None)], level=0))
-        module.body.insert(2, ast.ImportFrom(module='mindspore.nn', names=[ast.alias(name='Cell', asname=None)],
-                                             level=0))
-        module.body.insert(3, ast.ImportFrom(module='mindspore.ops', names=[ast.alias(name='functional', asname='F')],
-                                             level=0))
+        stree.get_import_asts().append(ast.Import([ast.alias(name='mindspore', asname=None)]))
+        stree.get_import_asts().append(ast.ImportFrom(module='mindspore', names=[ast.alias(name='nn', asname=None)],
+                                                      level=0))
+        stree.get_import_asts().append(ast.ImportFrom(module='mindspore.nn',
+                                                      names=[ast.alias(name='Cell', asname=None)], level=0))
+        stree.get_import_asts().append(ast.ImportFrom(module='mindspore.ops',
+                                                      names=[ast.alias(name='functional', asname='F')], level=0))
+        origin_net = stree.get_origin_network()
+        net_path = inspect.getfile(type(origin_net))
+        ModuleParser.save_file_path_to_sys(stree, 0, net_path)
+        ModuleParser.save_imports_from_file(stree, net_path)
+
+    @staticmethod
+    def get_valid_import_info(import_node, file_path):
+        """Get valid import info while import_node.module is at form of relative path"""
+        # copy to a new node to avoid origin import_node being modified.
+        import_node_test = copy.deepcopy(import_node)
+        file_path = os.path.dirname(os.path.abspath(file_path))
+        # get real path from import_node.level
+        # from .(A) import xxx: current path
+        # from ..(A) import xxx: last level path
+        import_node_module_name = import_node.module
+        level = import_node.level
+        # from A import xxx: it does not need to pad, directly return the module name
+        if level == 0:
+            return import_node_module_name, None
+        if level > 1:
+            for _ in range(level - 1):
+                file_path = os.path.dirname(file_path)
+        file_path_tmp = file_path[:]
+        max_level_count = file_path.count('/') + file_path.count('\\') - 1
+        level_count = 0
+        # suffix is the module_name, e.g. 'A' in 'from ..(A) import xxx'
+        suffix = ''
+        if import_node_module_name:
+            suffix = '.' + import_node_module_name
+        while level_count < max_level_count:
+            file_path_tmp = os.path.dirname(file_path_tmp)
+            import_node_test.module = file_path[len(file_path_tmp) + 1:].replace('/', '.') + suffix
+            import_node_test.level = 0
+            import_code = astunparse.unparse(import_node_test).strip()
+            test_code = f"import sys\nsys.path.insert(0, r'{file_path_tmp}')\n{import_code}"
+            try:
+                exec(test_code) # pylint: disable=W0122
+            except (ValueError, ImportError) as e:
+                # try upper level to avoid ValueError: attempted relative import beyond top-level package
+                # this exception is changed to ImportError after python3.9
+                logger.info(f"For MindSpore Rewrite, in module parser, test import code: "
+                            f"{import_code} failed: {e}. Try upper level.")
+                level_count += 1
+                continue
+            except Exception as e: # pylint: disable=W0703
+                logger.info(f"For MindSpore Rewrite, in module parser, process import code: "
+                            f"{import_code} failed: {e}. Ignore this import code.")
+                return None, None
+            else:
+                # try test code success
+                return import_node_test.module, file_path_tmp
+        # try codes with all level failed
+        logger.info(f"For MindSpore Rewrite, in module parser, test import code: "
+                    f"{astunparse.unparse(import_node).strip()} failed. Ignore this import code.")
+        return None, None
+
+    @staticmethod
+    def save_imports_from_file(stree, file_path):
+        """Save imports from file"""
+        if not os.path.exists(file_path):
+            raise RuntimeError(f"For MindSpore Rewrite, in module parser, file {file_path} not exist.")
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                source_code = f.read()
+                import_nodes = ModuleParser._get_import_node(ast.parse(source_code))
+        except RuntimeError as err:
+            raise RuntimeError(f"For MindSpore Rewrite, in module parser, get import nodes error: {err}")
+        if not import_nodes:
+            return
+        for import_node in import_nodes:
+            import_node = ModuleParser._process_relative_import(stree, import_node, file_path)
+            if import_node:
+                stree.get_import_asts().append(import_node)
+
+    @staticmethod
+    def _process_relative_import(stree, import_node, file_path):
+        """Process relative imports"""
+        if isinstance(import_node, ast.ImportFrom):
+            # pad the ImportFrom with parent path
+            # e.g. from ..C import xxx -> from A.B.C import xxx
+            import_module, import_path = ModuleParser.get_valid_import_info(import_node, file_path)
+            if import_path:
+                ModuleParser.save_file_path_to_sys(stree, 0, import_path)
+            module_name_list = [alias.name.strip() for alias in import_node.names]
+            # add the module into _imported_modules to direct the class
+            stree.save_imported_modules(file_path, import_module, module_name_list)
+            import_node = ast.ImportFrom(module=import_module, names=import_node.names, level=0)
+        elif isinstance(import_node, ast.Import):
+            for alias in import_node.names:
+                name = alias.name
+                stree.save_imported_modules(file_path, name.strip(), [])
+        return import_node
+
+    @staticmethod
+    def _add_decorator_to_class(class_ast: ast.ClassDef, origin_net):
+        """Add decorators to class"""
         origin_net_source_code_file = inspect.getfile(type(origin_net))
         if not os.path.exists(origin_net_source_code_file):
             raise RuntimeError("For MindSpore Rewrite, in module parser, File ", origin_net_source_code_file,
@@ -84,35 +228,62 @@ class ModuleParser(Parser):
         try:
             with open(origin_net_source_code_file, "r", encoding="utf-8") as f:
                 source_code = f.read()
-                import_nodes = ModuleParser._get_import_node(ast.parse(source_code))
+                decorators = ModuleParser._get_decorator(ast.parse(source_code), origin_net)
         except RuntimeError:
-            raise RuntimeError("For MindSpore Rewrite, in module parser, get import nodes error")
-        if import_nodes:
-            for import_index, import_node in enumerate(import_nodes):
-                module.body.insert(import_index + 4, import_node)
-        ast.fix_missing_locations(module)
+            raise RuntimeError("For MindSpore Rewrite, in module parser, get decorators error")
+        if decorators:
+            for decorator_index, decorator_node in enumerate(decorators):
+                class_ast.decorator_list.insert(decorator_index, decorator_node)
+        ast.fix_missing_locations(class_ast)
 
     @staticmethod
-    def _save_net_file_path(stree: SymbolTree):
-        origin_net_file = inspect.getfile(type(stree.get_origin_network()))
-        file_full_path = os.path.abspath(origin_net_file)
-        file_path, _ = os.path.split(file_full_path)
-        stree.append_net_file_path(file_path)
+    def _get_decorator(ast_root, origin_net):
+        """Get the decorators of function"""
+        net_name = type(origin_net).__name__
+        decorators = []
+
+        class GetClassNode(ast.NodeVisitor):
+            """Find the class node from input ast node."""
+            def visit_ClassDef(self, node: ast.ClassDef) -> Any:
+                """Visit the class node and add the decorators to class node"""
+                if node.name == net_name:
+                    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
+                        # User should set the denied class_decorator,
+                        # because the symbol_tree cant pass the correct parameters to decorators but the instance "obj".
+                        if decorator_name not in ModuleParser.denied_class_decorator_list:
+                            decorators.append(decorator)
+                return node
+
+            def get_node(self, input_ast):
+                """Interface of GetClassNode."""
+                self.generic_visit(input_ast)
+                return True
+
+        get_node_handler = GetClassNode()
+        get_node_handler.get_node(ast_root)
+        return decorators
 
     def target(self):
         """Parse target type"""
         return ast.Module
 
-    def process(self, stree: SymbolTree, node: ast.Module):
+    def process(self, stree: SymbolTree, node: ast.Module, node_manager: NodeManager):
         """Process ast.ClassDef nodes in ast.Module."""
-        ModuleParser._add_import_to_module(node, stree.get_origin_network())
+        ModuleParser._save_imports(stree)
         class_ast = ModuleParser._find_class(node)
+        ModuleParser._add_decorator_to_class(class_ast, stree.get_origin_network())
         stree.set_class_ast(class_ast)
-        ModuleParser._save_net_file_path(stree)
         for body in node.body:
             if isinstance(body, ast.ClassDef):
                 parser: Parser = ParserRegister.instance().get_parser(ast.ClassDef)
-                parser.process(stree, body)
+                parser.process(stree, body, stree)
             else:
                 logger.info(f"For MindSpore Rewrite, in module parser, Ignoring unsupported "
                             f"node({astunparse.unparse(body)}) in ast.Module.")

mindspore/rewrite/{parser.py → parsers/parser.py}

@@ -16,7 +16,8 @@
 import abc
 import ast
 
-from .symbol_tree import SymbolTree
+from ..symbol_tree import SymbolTree
+from ..node.node_manager import NodeManager
 
 
 class Parser(abc.ABC):
@@ -34,12 +35,13 @@ class Parser(abc.ABC):
         return type(None)
 
     @abc.abstractmethod
-    def process(self, stree: SymbolTree, node: ast.AST):
+    def process(self, stree: SymbolTree, node: ast.AST, node_manager: NodeManager):
         """
         Parse input ast node and add parse result into SymbolTree.
 
         Args:
              stree (SymbolTree): current symbol_tree
              node (ast.AST): node who is tried to be parsed
+             node_manager (NodeManager): NodeManager those asts belong to.
         """
         raise NotImplementedError

mindspore/rewrite/{parser_register.py → parsers/parser_register.py}

@@ -45,7 +45,7 @@ class ParserRegister:
             parser (Parser): An instance of Parser to be registered.
         """
         if isinstance(parser, Parser):
-            ParserRegister.instance()._parsers[parser.target()] = parser
+            ParserRegister.instance().get_parsers()[parser.target()] = parser
 
     def get_parser(self, ast_type: type) -> Optional[Parser]:
         """

mindspore/rewrite/parsers/return_parser.py

@@ -16,9 +16,10 @@
 import ast
 
 from ..symbol_tree import SymbolTree
-from ..node import Node
-from ..parser import Parser
-from ..parser_register import reg_parser
+from ..node.node import Node
+from ..node.node_manager import NodeManager
+from .parser import Parser
+from .parser_register import reg_parser
 from ..common import error_str
 
 
@@ -29,14 +30,13 @@ class ReturnParser(Parser):
         """Parse target type"""
         return ast.Return
 
-    def process(self, stree: SymbolTree, node: ast.Return):
+    def process(self, stree: SymbolTree, node: ast.Return, node_manager: NodeManager):
         """Parse ast.Return to output-node of SymbolTree."""
         return_value = node.value
         if not isinstance(return_value, ast.Name):
             raise RuntimeError(error_str(f"only support ast.Name as return value, but got ast type "
                                          f"'{type(return_value).__name__}'", father_node=node, child_node=return_value))
         node_return = Node.create_output_node(node, [return_value.id])
-        stree.append_origin_field(node_return)
-
+        stree.append_origin_field(node_return, node_manager)
 
 g_return_parser = reg_parser(ReturnParser())

mindspore/rewrite/sparsify/sparse_transformer.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 # ============================================================================
 """Sparsify transformer"""
+import sys
 import ast
 import inspect
 import textwrap
 from collections import deque
-import astunparse
 
 from mindspore import ops, nn
 from mindspore import log as logger
@@ -25,6 +25,10 @@ from mindspore.rewrite.parsers.assign_parser import AssignParser
 from mindspore.rewrite.sparsify.utils import ArgType, SparseFunc, sparse_rules, get_sparse_func, builtin_ops, \
     get_binop_name, get_sparse_method_outputs, arg_type_to_prefix_map, get_inputs_outputs
 
+if sys.version_info >= (3, 9):
+    import ast as astunparse # pylint: disable=reimported, ungrouped-imports
+else:
+    import astunparse
 
 OPS_MODULE = "mindspore.ops."
 MAX_RECURSION_DEPTH = 10
@@ -61,8 +65,13 @@ def sparsify_helper(f, arg_types, user_defined_rules=None, sparse_name="", full_
 
     if changed:
         sparse_tree = list(x[0] for x in sparse_transformer.sparse_functiondef.values()) + sparse_tree
-        ast_module = ast.Module([ast.FunctionDef(
-            sparse_name, functiondef.args, sparse_tree, functiondef.decorator_list, functiondef.returns)])
+        if sys.version_info >= (3, 9):
+            ast_module = ast.Module([ast.FunctionDef(
+                sparse_name, functiondef.args, sparse_tree, functiondef.decorator_list, functiondef.returns)],
+                                    type_ignores=[])
+        else:
+            ast_module = ast.Module([ast.FunctionDef(
+                sparse_name, functiondef.args, sparse_tree, functiondef.decorator_list, functiondef.returns)])
         return ast_module, True, return_types
     return tree, False, return_types
 

mindspore/rewrite/sparsify/utils.py

@@ -176,4 +176,4 @@ def get_binop_name(binop):
         return "*"
     if binop == ast.Div():
         return "/"
-    return None
+    return ""