-
-
Notifications
You must be signed in to change notification settings - Fork 0
/
no_assert_codemod.py
46 lines (42 loc) · 1.48 KB
/
no_assert_codemod.py
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
import libcst as cst
from libcst.codemod import VisitorBasedCodemodCommand
def _get_assert_replacement(node: cst.Assert):
if node.msg:
message = node.msg
else:
message = cst.SimpleString(value=repr(str(cst.Module(body=[node]).code)))
return cst.If(
test=cst.UnaryOperation(
operator=cst.Not(),
expression=node.test, # Todo: parenthesize?
),
body=cst.IndentedBlock(
body=[
cst.SimpleStatementLine(
body=[
cst.Raise(
exc=cst.Call(
func=cst.Name(
value="AssertionError",
),
args=[
cst.Arg(value=message),
],
),
),
]
),
],
),
)
class RewriteAssertToIfRaise(VisitorBasedCodemodCommand):
def leave_SimpleStatementLine(
self,
original_node: cst.SimpleStatementLine,
updated_node: cst.SimpleStatementLine,
) -> cst.SimpleStatementLine:
if isinstance(updated_node.body[0], cst.Assert):
if_stmt = _get_assert_replacement(updated_node.body[0])
replacement = updated_node.deep_replace(updated_node, if_stmt)
return replacement
return updated_node