1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from textwrap import dedent
from typing import Callable
import libcst as cst
from libcst.metadata import ExperimentalReentrantCodegenProvider, MetadataWrapper
from libcst.testing.utils import data_provider, UnitTest
class ExperimentalReentrantCodegenProviderTest(UnitTest):
@data_provider(
{
"simple_top_level_statement": {
"old_module": (
"""\
import math
c = math.sqrt(a*a + b*b)
"""
),
"new_module": (
"""\
import math
c = math.hypot(a, b)
"""
),
"old_node": lambda m: m.body[1],
"new_node": cst.parse_statement("c = math.hypot(a, b)"),
},
"replacement_inside_block": {
"old_module": (
"""\
import math
def do_math(a, b):
c = math.sqrt(a*a + b*b)
return c
"""
),
"new_module": (
"""\
import math
def do_math(a, b):
c = math.hypot(a, b)
return c
"""
),
"old_node": lambda m: m.body[1].body.body[0],
"new_node": cst.parse_statement("c = math.hypot(a, b)"),
},
"missing_trailing_newline": {
"old_module": "old_fn()", # this module has no trailing newline
"new_module": "new_fn()",
"old_node": lambda m: m.body[0],
"new_node": cst.parse_statement("new_fn()\n"),
},
"nested_blocks_with_missing_trailing_newline": {
"old_module": (
"""\
if outer:
if inner:
old_fn()""" # this module has no trailing newline
),
"new_module": (
"""\
if outer:
if inner:
new_fn()"""
),
"old_node": lambda m: m.body[0].body.body[0].body.body[0],
"new_node": cst.parse_statement("new_fn()\n"),
},
}
)
def test_provider(
self,
old_module: str,
new_module: str,
old_node: Callable[[cst.Module], cst.CSTNode],
new_node: cst.BaseStatement,
) -> None:
old_module = dedent(old_module)
new_module = dedent(new_module)
mw = MetadataWrapper(cst.parse_module(old_module))
codegen_partial = mw.resolve(ExperimentalReentrantCodegenProvider)[
old_node(mw.module)
]
self.assertEqual(codegen_partial.get_original_module_code(), old_module)
self.assertEqual(codegen_partial.get_modified_module_code(new_node), new_module)
def test_byte_conversion(
self,
) -> None:
module_bytes = "fn()\n".encode("utf-16")
mw = MetadataWrapper(
cst.parse_module("fn()\n", cst.PartialParserConfig(encoding="utf-16"))
)
codegen_partial = mw.resolve(ExperimentalReentrantCodegenProvider)[
mw.module.body[0]
]
self.assertEqual(codegen_partial.get_original_module_bytes(), module_bytes)
self.assertEqual(
codegen_partial.get_modified_module_bytes(cst.parse_statement("fn2()\n")),
"fn2()\n".encode("utf-16"),
)
|
