mirror of
https://github.com/Instagram/LibCST.git
synced 2025-12-23 10:35:53 +00:00
112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
# 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 Dict
|
|
|
|
import libcst as cst
|
|
from libcst.codemod import (
|
|
Codemod,
|
|
CodemodContext,
|
|
CodemodTest,
|
|
SkipFile,
|
|
transform_module,
|
|
TransformExit,
|
|
TransformFailure,
|
|
TransformSkip,
|
|
TransformSuccess,
|
|
)
|
|
|
|
|
|
class TestRunner(CodemodTest):
|
|
def test_runner_default(self) -> None:
|
|
before = """
|
|
def foo() -> None:
|
|
pass
|
|
|
|
def bar() -> int:
|
|
return 5
|
|
"""
|
|
after = """
|
|
# A comment
|
|
def foo() -> None:
|
|
pass
|
|
|
|
def bar() -> int:
|
|
return 5
|
|
"""
|
|
|
|
class SimpleCodemod(Codemod):
|
|
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
|
|
self.warn("Testing")
|
|
return tree.with_changes(
|
|
header=[cst.EmptyLine(comment=cst.Comment("# A comment"))]
|
|
)
|
|
|
|
transform = SimpleCodemod(CodemodContext())
|
|
response = transform_module(transform, dedent(before))
|
|
self.assertIsInstance(response, TransformSuccess)
|
|
assert isinstance(response, TransformSuccess)
|
|
self.assertCodeEqual(response.code, after)
|
|
self.assertEqual(response.warning_messages, ["Testing"])
|
|
|
|
def test_runner_interrupted(self) -> None:
|
|
code = """
|
|
def foo() -> None:
|
|
pass
|
|
|
|
def bar() -> int:
|
|
return 5
|
|
"""
|
|
|
|
class SimpleCodemod(Codemod):
|
|
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
|
|
raise KeyboardInterrupt("Testing")
|
|
|
|
transform = SimpleCodemod(CodemodContext())
|
|
response = transform_module(transform, dedent(code))
|
|
self.assertIsInstance(response, TransformExit)
|
|
|
|
def test_runner_skip(self) -> None:
|
|
code = """
|
|
def foo() -> None:
|
|
pass
|
|
|
|
def bar() -> int:
|
|
return 5
|
|
"""
|
|
|
|
class SimpleCodemod(Codemod):
|
|
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
|
|
self.warn("Testing")
|
|
raise SkipFile()
|
|
|
|
transform = SimpleCodemod(CodemodContext())
|
|
response = transform_module(transform, dedent(code))
|
|
self.assertIsInstance(response, TransformSkip)
|
|
self.assertEqual(response.warning_messages, ["Testing"])
|
|
|
|
def test_runner_failure(self) -> None:
|
|
code = """
|
|
def foo() -> None:
|
|
pass
|
|
|
|
def bar() -> int:
|
|
return 5
|
|
"""
|
|
|
|
class SimpleCodemod(Codemod):
|
|
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
|
|
self.warn("Testing")
|
|
somedict: Dict[str, str] = {}
|
|
somedict["invalid_key"]
|
|
return tree
|
|
|
|
transform = SimpleCodemod(CodemodContext())
|
|
response = transform_module(transform, dedent(code))
|
|
self.assertIsInstance(response, TransformFailure)
|
|
assert isinstance(response, TransformFailure)
|
|
self.assertEqual(response.warning_messages, ["Testing"])
|
|
self.assertIsInstance(response.error, KeyError)
|