[Dy2St] Move transformers to a separate dir (#60246)

SigureMo · web-flow · commit 5faf595c036d · 2023-12-25T10:03:29.000+08:00
diff --git a/python/paddle/jit/dy2static/__init__.py b/python/paddle/jit/dy2static/__init__.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .assert_transformer import AssertTransformer  # noqa: F401
-from .ast_transformer import DygraphToStaticAst  # noqa: F401
 from .convert_call_func import convert_call as Call  # noqa: F401
 from .convert_operators import (  # noqa: F401
     convert_assert as Assert,
@@ -33,6 +31,7 @@
 )
 from .program_translator import convert_to_static  # noqa: F401
 from .static_analysis import NodeVarType, StaticAnalysisVisitor  # noqa: F401
+from .transformers import DygraphToStaticAst  # noqa: F401
 from .utils import UndefinedVar, ast_to_source_code, saw  # noqa: F401
 from .variable_trans_func import (  # noqa: F401
     create_bool_as_type,
diff --git a/python/paddle/jit/dy2static/ast_utils.py b/python/paddle/jit/dy2static/ast_utils.py
@@ -14,6 +14,7 @@
 
 
 import ast
+import sys
 
 import astor
 
@@ -32,6 +33,10 @@ def ast_to_source_code(ast_node):
     if isinstance(ast_node, gast.AST):
         ast_node = gast.gast_to_ast(ast_node)
 
+    if sys.version_info >= (3, 9):
+        ast.fix_missing_locations(ast_node)
+        return ast.unparse(ast_node)
+
     # Do not wrap lines even if they are too long
     def pretty_source(source):
         return ''.join(source)
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -37,7 +37,6 @@
 from paddle.utils import flatten, gast
 
 from . import error, logging_utils
-from .ast_transformer import DygraphToStaticAst
 from .function_spec import (
     FunctionSpec,
     _hash_spec_names,
@@ -53,6 +52,7 @@
 from .pir_partial_program import (
     PartialProgramLayerHook as PirPartialProgramLayerHook,
 )
+from .transformers import DygraphToStaticAst
 from .utils import (
     ALREADY_D2S,
     NO_SHAPE_VAR_TYPE,
diff --git a/python/paddle/jit/dy2static/transformers/__init__.py b/python/paddle/jit/dy2static/transformers/__init__.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .transform import DygraphToStaticAst  # noqa: F401
diff --git a/python/paddle/jit/dy2static/transformers/assert_transformer.py b/python/paddle/jit/dy2static/transformers/assert_transformer.py
@@ -15,7 +15,7 @@
 from paddle.jit.dy2static.utils import ast_to_source_code
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/base.py b/python/paddle/jit/dy2static/transformers/base.py
@@ -23,7 +23,6 @@
     ORIGI_INFO,
     ast_to_source_code,
     create_assign_node,
-    get_attribute_full_name,
 )
 from paddle.utils import gast
 
@@ -50,35 +49,6 @@ def visit(self, node):
         return result
 
 
-class RenameTransformer(BaseTransformer):
-    def __init__(self, node):
-        assert isinstance(
-            node, gast.AST
-        ), "RenameTransformer only accepts gast.AST as input"
-        self.root = node
-        self.old_name = ""
-        self.new_name = ""
-
-    def rename(self, old_name, new_name):
-        self.old_name = old_name
-        self.new_name = new_name
-        self.visit(self.root)
-
-    def visit_Name(self, node):
-        self.generic_visit(node)
-        if node.id == self.old_name:
-            node.id = self.new_name
-        return node
-
-    def visit_Attribute(self, node):
-        self.generic_visit(node)
-        attr_full_name = get_attribute_full_name(node)
-        if attr_full_name == self.old_name:
-            new_name_node = gast.parse(self.new_name).body[0].value
-            return new_name_node
-        return node
-
-
 class NameNodeReplaceTransformer(BaseTransformer):
     """
     This class replaces specified gast.Name node by replace_node.
diff --git a/python/paddle/jit/dy2static/transformers/basic_api_transformer.py b/python/paddle/jit/dy2static/transformers/basic_api_transformer.py
@@ -17,8 +17,8 @@
 
 from paddle.utils import gast
 
-from . import utils
-from .base_transformer import BaseTransformer
+from .. import utils
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/break_continue_transformer.py b/python/paddle/jit/dy2static/transformers/break_continue_transformer.py
@@ -17,7 +17,7 @@
 from paddle.jit.dy2static.variable_trans_func import create_bool_node
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer, ForNodeVisitor
+from .base import BaseTransformer, ForNodeVisitor
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/call_transformer.py b/python/paddle/jit/dy2static/transformers/call_transformer.py
@@ -15,8 +15,8 @@
 from paddle.jit.dy2static.utils import ast_to_source_code, is_paddle_api
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
-from .utils import is_builtin  # noqa: F401
+from ..utils import is_builtin  # noqa: F401
+from .base import BaseTransformer
 
 PDB_SET = "pdb.set_trace"
 
diff --git a/python/paddle/jit/dy2static/transformers/cast_transformer.py b/python/paddle/jit/dy2static/transformers/cast_transformer.py
@@ -15,7 +15,7 @@
 from paddle.jit.dy2static.utils import ast_to_source_code
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/create_variable_transformer.py b/python/paddle/jit/dy2static/transformers/create_variable_transformer.py
@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base_transformer import BaseTransformer
-from .utils import FunctionNameLivenessAnalysis
-from .variable_trans_func import create_undefined_var
+from ..utils import FunctionNameLivenessAnalysis
+from ..variable_trans_func import create_undefined_var
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/decorator_transformer.py b/python/paddle/jit/dy2static/transformers/decorator_transformer.py
@@ -18,8 +18,8 @@
 
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
-from .utils import RE_PYMODULE, RE_PYNAME, ast_to_source_code
+from ..utils import RE_PYMODULE, RE_PYNAME, ast_to_source_code
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/early_return_transformer.py b/python/paddle/jit/dy2static/transformers/early_return_transformer.py
@@ -14,7 +14,7 @@
 
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/ifelse_transformer.py b/python/paddle/jit/dy2static/transformers/ifelse_transformer.py
@@ -41,8 +41,8 @@
 # See details in https://github.com/serge-sans-paille/gast/
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
-from .utils import FALSE_FUNC_PREFIX, TRUE_FUNC_PREFIX
+from ..utils import FALSE_FUNC_PREFIX, TRUE_FUNC_PREFIX
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/logical_transformer.py b/python/paddle/jit/dy2static/transformers/logical_transformer.py
@@ -14,8 +14,8 @@
 
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
-from .utils import ast_to_source_code
+from ..utils import ast_to_source_code
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/loop_transformer.py b/python/paddle/jit/dy2static/transformers/loop_transformer.py
@@ -18,14 +18,8 @@
 from paddle.base import unique_name
 from paddle.utils import gast
 
-from .base_transformer import (
-    BaseTransformer,
-    ForLoopTuplePreTransformer,
-    ForNodeVisitor,
-)
-from .ifelse_transformer import ARGS_NAME
-from .static_analysis import NodeVarType, StaticAnalysisVisitor
-from .utils import (
+from ..static_analysis import NodeVarType, StaticAnalysisVisitor
+from ..utils import (
     FOR_BODY_PREFIX,
     FOR_CONDITION_PREFIX,
     WHILE_BODY_PREFIX,
@@ -39,6 +33,12 @@
     create_set_args_node,
     get_attribute_full_name,
 )
+from .base import (
+    BaseTransformer,
+    ForLoopTuplePreTransformer,
+    ForNodeVisitor,
+)
+from .ifelse_transformer import ARGS_NAME
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/return_transformer.py b/python/paddle/jit/dy2static/transformers/return_transformer.py
@@ -15,14 +15,14 @@
 from paddle.base import unique_name
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
-from .break_continue_transformer import ForToWhileTransformer
-from .utils import (
+from ..utils import (
     ORIGI_INFO,
     Dygraph2StaticException,
     ast_to_source_code,
     index_in_list,
 )
+from .base import BaseTransformer
+from .break_continue_transformer import ForToWhileTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/tensor_shape_transformer.py b/python/paddle/jit/dy2static/transformers/tensor_shape_transformer.py
@@ -14,8 +14,8 @@
 
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
-from .utils import ast_to_source_code
+from ..utils import ast_to_source_code
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/tensorhook_transformer.py b/python/paddle/jit/dy2static/transformers/tensorhook_transformer.py
@@ -16,7 +16,7 @@
 
 from paddle.utils import gast
 
-from .base_transformer import BaseTransformer
+from .base import BaseTransformer
 
 
 class RegisterHookTransformer(BaseTransformer):
diff --git a/python/paddle/jit/dy2static/transformers/transform.py b/python/paddle/jit/dy2static/transformers/transform.py
@@ -19,9 +19,10 @@
 
 import os
 
-from . import logging_utils
+from .. import logging_utils
+from ..utils import ast_to_source_code
 from .assert_transformer import AssertTransformer
-from .base_transformer import BaseTransformer
+from .base import BaseTransformer
 from .basic_api_transformer import BasicApiTransformer, NameloadJstTransformer
 from .break_continue_transformer import (
     BreakContinueTransformer,
@@ -39,7 +40,6 @@
 from .tensor_shape_transformer import TensorShapeTransformer
 from .tensorhook_transformer import RegisterHookTransformer
 from .typehint_transformer import TypeHintTransformer
-from .utils import ast_to_source_code
 
 __all__ = []
 
diff --git a/python/paddle/jit/dy2static/transformers/typehint_transformer.py b/python/paddle/jit/dy2static/transformers/typehint_transformer.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from .base_transformer import BaseTransformer
+from .base import BaseTransformer
 
 __all__ = []
 
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
@@ -1611,7 +1611,9 @@ def expand_undefined_var(nest1, nest2, names):
     nest2: Var2, ([1,2,3,4], UndefinedVar)
     In this case, we should not expand recursively.
     """
-    from paddle.jit.dy2static.return_transformer import RETURN_VALUE_PREFIX
+    from paddle.jit.dy2static.transformers.return_transformer import (
+        RETURN_VALUE_PREFIX,
+    )
     from paddle.jit.dy2static.utils import UndefinedVar
 
     def pack_undefined_var_as(seq):
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -462,6 +462,7 @@ packages=['paddle',
           'paddle.framework',
           'paddle.jit',
           'paddle.jit.dy2static',
+          'paddle.jit.dy2static.transformers',
           'paddle.jit.pir_dy2static',
           'paddle.jit.sot',
           'paddle.jit.sot.opcode_translator',
diff --git a/setup.py b/setup.py
@@ -1465,6 +1465,7 @@ def get_setup_parameters():
         'paddle.framework',
         'paddle.jit',
         'paddle.jit.dy2static',
+        'paddle.jit.dy2static.transformers',
         'paddle.jit.pir_dy2static',
         'paddle.jit.sot',
         'paddle.jit.sot.opcode_translator',
diff --git a/test/dygraph_to_static/test_logical.py b/test/dygraph_to_static/test_logical.py
@@ -21,7 +21,9 @@
 from dygraph_to_static_utils import Dy2StTestBase, enable_to_static_guard
 
 import paddle
-from paddle.jit.dy2static.logical_transformer import cmpop_node_to_str
+from paddle.jit.dy2static.transformers.logical_transformer import (
+    cmpop_node_to_str,
+)
 from paddle.utils import gast
 
 SEED = 2020
diff --git a/test/dygraph_to_static/test_loop.py b/test/dygraph_to_static/test_loop.py
@@ -26,7 +26,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle import base
-from paddle.jit.dy2static.loop_transformer import NameVisitor
+from paddle.jit.dy2static.transformers.loop_transformer import NameVisitor
 from paddle.utils import gast
 
 SEED = 2020
