[libclang/python] Expose Rewriter to the python binding (#77269)

Exposes `CXRewriter` API to the python binding as `class Rewriter`.

2 files changed, 137 insertions, 0 deletions

diff --git a/clang/bindings/python/clang/cindex.py b/clang/bindings/python/clang/cindex.py
index 754f03d..44a34ca 100644
--- a/clang/bindings/python/clang/cindex.py
+++ b/clang/bindings/python/clang/cindex.py
@@ -3506,6 +3506,65 @@ class Token(Structure):
         return cursor
 
 
+class Rewriter(ClangObject):
+    """
+    The Rewriter is a wrapper class around clang::Rewriter
+
+    It enables rewriting buffers.
+    """
+
+    @staticmethod
+    def create(tu):
+        """
+        Creates a new Rewriter
+        Parameters:
+        tu -- The translation unit for the target AST.
+        """
+        return Rewriter(conf.lib.clang_CXRewriter_create(tu))
+
+    def __init__(self, ptr):
+        ClangObject.__init__(self, ptr)
+
+    def __del__(self):
+        conf.lib.clang_CXRewriter_dispose(self)
+
+    def insert_text_before(self, loc, insert):
+        """
+        Insert the specified string at the specified location in
+        the original buffer.
+        """
+        conf.lib.clang_CXRewriter_insertTextBefore(self, loc, insert)
+
+    def replace_text(self, extent, replacement):
+        """
+        This method replaces a range of characters in the input buffer with
+        a new string.
+        """
+        conf.lib.clang_CXRewriter_replaceText(self, extent, replacement)
+
+    def remove_text(self, extent):
+        """
+        Remove the specified text region.
+        """
+        conf.lib.clang_CXRewriter_removeText(self, extent)
+
+    def overwrite_changed_files(self):
+        """
+        Save all changed files to disk.
+
+        Returns 1 if any files were not saved successfully,
+        returns 0 otherwise.
+        """
+        return conf.lib.clang_CXRewriter_overwriteChangedFiles(self)
+
+    def write_main_file_to_stdout(self):
+        """
+        Writes the main file to stdout.
+        """
+        sys.stdout.flush()
+        conf.lib.clang_CXRewriter_writeMainFileToStdOut(self)
+
+
 # Now comes the plumbing to hook up the C library.
 
 # Register callback types in common container.
@@ -3571,6 +3630,13 @@ functionList = [
     ("clang_codeCompleteGetNumDiagnostics", [CodeCompletionResults], c_int),
     ("clang_createIndex", [c_int, c_int], c_object_p),
     ("clang_createTranslationUnit", [Index, c_interop_string], c_object_p),
+    ("clang_CXRewriter_create", [TranslationUnit], c_object_p),
+    ("clang_CXRewriter_dispose", [Rewriter]),
+    ("clang_CXRewriter_insertTextBefore", [Rewriter, SourceLocation, c_interop_string]),
+    ("clang_CXRewriter_overwriteChangedFiles", [Rewriter], c_int),
+    ("clang_CXRewriter_removeText", [Rewriter, SourceRange]),
+    ("clang_CXRewriter_replaceText", [Rewriter, SourceRange, c_interop_string]),
+    ("clang_CXRewriter_writeMainFileToStdOut", [Rewriter]),
     ("clang_CXXConstructor_isConvertingConstructor", [Cursor], bool),
     ("clang_CXXConstructor_isCopyConstructor", [Cursor], bool),
     ("clang_CXXConstructor_isDefaultConstructor", [Cursor], bool),
diff --git a/clang/bindings/python/tests/cindex/test_rewrite.py b/clang/bindings/python/tests/cindex/test_rewrite.py
new file mode 100644
index 0000000..42006f5
--- /dev/null
+++ b/clang/bindings/python/tests/cindex/test_rewrite.py
@@ -0,0 +1,71 @@
+import unittest
+import tempfile
+
+from clang.cindex import (
+    Rewriter,
+    TranslationUnit,
+    File,
+    SourceLocation,
+    SourceRange,
+)
+
+
+class TestRewrite(unittest.TestCase):
+    code = """int main() { return 0; }"""
+
+    def setUp(self):
+        self.tmp = tempfile.NamedTemporaryFile(suffix=".cpp", buffering=0)
+        self.tmp.write(TestRewrite.code.encode("utf-8"))
+        self.tmp.flush()
+        self.tu = TranslationUnit.from_source(self.tmp.name)
+        self.rew = Rewriter.create(self.tu)
+        self.file = File.from_name(self.tu, self.tmp.name)
+
+    def tearDown(self):
+        self.tmp.close()
+
+    def get_content(self) -> str:
+        with open(self.tmp.name, "r", encoding="utf-8") as f:
+            return f.read()
+
+    def test_replace(self):
+        rng = SourceRange.from_locations(
+            SourceLocation.from_position(self.tu, self.file, 1, 5),
+            SourceLocation.from_position(self.tu, self.file, 1, 9),
+        )
+        self.rew.replace_text(rng, "MAIN")
+        self.rew.overwrite_changed_files()
+        self.assertEqual(self.get_content(), "int MAIN() { return 0; }")
+
+    def test_replace_shorter(self):
+        rng = SourceRange.from_locations(
+            SourceLocation.from_position(self.tu, self.file, 1, 5),
+            SourceLocation.from_position(self.tu, self.file, 1, 9),
+        )
+        self.rew.replace_text(rng, "foo")
+        self.rew.overwrite_changed_files()
+        self.assertEqual(self.get_content(), "int foo() { return 0; }")
+
+    def test_replace_longer(self):
+        rng = SourceRange.from_locations(
+            SourceLocation.from_position(self.tu, self.file, 1, 5),
+            SourceLocation.from_position(self.tu, self.file, 1, 9),
+        )
+        self.rew.replace_text(rng, "patatino")
+        self.rew.overwrite_changed_files()
+        self.assertEqual(self.get_content(), "int patatino() { return 0; }")
+
+    def test_insert(self):
+        pos = SourceLocation.from_position(self.tu, self.file, 1, 5)
+        self.rew.insert_text_before(pos, "ro")
+        self.rew.overwrite_changed_files()
+        self.assertEqual(self.get_content(), "int romain() { return 0; }")
+
+    def test_remove(self):
+        rng = SourceRange.from_locations(
+            SourceLocation.from_position(self.tu, self.file, 1, 5),
+            SourceLocation.from_position(self.tu, self.file, 1, 9),
+        )
+        self.rew.remove_text(rng)
+        self.rew.overwrite_changed_files()
+        self.assertEqual(self.get_content(), "int () { return 0; }")