From eb062cc504027e1cebdc214c843a221c30244e25 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Wed, 21 Jun 2023 21:17:12 -0700
Subject: [PATCH] refac and func update_files

---
 aider/coders/base_coder.py      |  8 ++++----
 aider/coders/editblock_coder.py |  4 +++-
 aider/coders/func_coder.py      | 26 ++++++++++++++++++++++++--
 aider/coders/wholefile_coder.py |  7 ++++---
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py
index d7315fb04..004d96fae 100755
--- a/aider/coders/base_coder.py
+++ b/aider/coders/base_coder.py
@@ -377,7 +377,7 @@ class Coder:
             self.cur_messages += [dict(role="assistant", content=content)]
             return
 
-        edited, edit_error = self.apply_updates(content)
+        edited, edit_error = self.apply_updates()
         if edit_error:
             return edit_error
 
@@ -544,7 +544,7 @@ class Coder:
         md = Markdown(show_resp, style=self.assistant_output_color, code_theme="default")
         live.update(md)
 
-    def modify_incremental_response(self):
+    def modify_incremental_response(self, final):
         return self.partial_response_content
 
     def get_context_from_history(self, history):
@@ -716,9 +716,9 @@ class Coder:
     def get_addable_relative_files(self):
         return set(self.get_all_relative_files()) - set(self.get_inchat_relative_files())
 
-    def apply_updates(self, content):
+    def apply_updates(self):
         try:
-            edited = self.update_files(content)
+            edited = self.update_files()
             return edited, None
         except ValueError as err:
             err = err.args[0]
diff --git a/aider/coders/editblock_coder.py b/aider/coders/editblock_coder.py
index 575779b87..c829175da 100644
--- a/aider/coders/editblock_coder.py
+++ b/aider/coders/editblock_coder.py
@@ -16,7 +16,9 @@ class EditBlockCoder(Coder):
     def update_cur_messages(self, content, edited):
         self.cur_messages += [dict(role="assistant", content=content)]
 
-    def update_files(self, content):
+    def update_files(self):
+        content = self.partial_response_content
+
         # might raise ValueError for malformed ORIG/UPD blocks
         edits = list(find_original_update_blocks(content))
 
diff --git a/aider/coders/func_coder.py b/aider/coders/func_coder.py
index 87fa39025..9673d629e 100644
--- a/aider/coders/func_coder.py
+++ b/aider/coders/func_coder.py
@@ -1,4 +1,5 @@
 import os
+from pathlib import Path
 
 from aider import diffs
 
@@ -102,5 +103,26 @@ class FunctionCoder(Coder):
 
         return "\n".join(show_diff)
 
-    def update_files(self, content):
-        pass
+    def update_files(self):
+        args = self.parse_partial_args()
+
+        if not args:
+            return
+
+        files = args.get("files", [])
+
+        edited = set()
+        for file_upd in files:
+            path = file_upd.get("path")
+            if not path:
+                raise ValueError(f"Missing path: {file_upd}")
+
+            content = file_upd.get("content")
+            if not content:
+                raise ValueError(f"Missing content: {file_upd}")
+
+            full_path = os.path.abspath(os.path.join(self.root, path))
+            Path(full_path).write_text(content)
+            edited.add(path)
+
+        return edited
diff --git a/aider/coders/wholefile_coder.py b/aider/coders/wholefile_coder.py
index da870f8ca..f9597f876 100644
--- a/aider/coders/wholefile_coder.py
+++ b/aider/coders/wholefile_coder.py
@@ -21,10 +21,11 @@ class WholeFileCoder(Coder):
             self.cur_messages += [dict(role="assistant", content=content)]
 
     def modify_incremental_response(self, final):
-        resp = self.partial_response_content
-        return self.update_files(resp, mode="diff")
+        return self.update_files(mode="diff")
+
+    def update_files(self, mode="update"):
+        content = self.partial_response_content
 
-    def update_files(self, content, mode="update"):
         edited = set()
         chat_files = self.get_inchat_relative_files()
         if not chat_files: