From 9b3db78b95b1a185e3cf560ed1c6bd63bc8188d3 Mon Sep 17 00:00:00 2001
From: David Cameron <davejcameron@users.noreply.github.com>
Date: Tue, 20 Jun 2023 06:09:28 -0400
Subject: [PATCH] Additional filename handling (#221)

---
 gpt_engineer/chat_to_files.py |  6 ++++
 tests/test_chat_to_files.py   | 52 +++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/gpt_engineer/chat_to_files.py b/gpt_engineer/chat_to_files.py
index f4812d4..e8c8270 100644
--- a/gpt_engineer/chat_to_files.py
+++ b/gpt_engineer/chat_to_files.py
@@ -14,6 +14,12 @@ def parse_chat(chat):  # -> List[Tuple[str, str]]:
         # Remove leading and trailing brackets
         path = re.sub(r"^\[(.*)\]$", r"\1", path)
 
+        # Remove leading and trailing backticks
+        path = re.sub(r"^`(.*)`$", r"\1", path)
+
+        # Remove trailing ]
+        path = re.sub(r"\]$", "", path)
+
         # Get the code
         code = match.group(2)
 
diff --git a/tests/test_chat_to_files.py b/tests/test_chat_to_files.py
index 176c19b..8c40bfd 100644
--- a/tests/test_chat_to_files.py
+++ b/tests/test_chat_to_files.py
@@ -92,3 +92,55 @@ def test_files_with_brackets_in_name():
 
     for file_name, file_content in expected_files.items():
         assert workspace[file_name] == file_content
+
+
+def test_files_with_file_colon():
+    chat = textwrap.dedent(
+        """
+    This is a sample program.
+
+    [FILE: file1.py]
+    ```python
+    print("Hello, World!")
+    ```
+    """
+    )
+
+    workspace = {}
+    to_files(chat, workspace)
+
+    assert workspace["all_output.txt"] == chat
+
+    expected_files = {
+        "file1.py": 'print("Hello, World!")\n',
+        "README.md": "\nThis is a sample program.\n\n[FILE: file1.py]\n",
+    }
+
+    for file_name, file_content in expected_files.items():
+        assert workspace[file_name] == file_content
+
+
+def test_files_with_back_tick():
+    chat = textwrap.dedent(
+        """
+    This is a sample program.
+
+    `file1.py`
+    ```python
+    print("Hello, World!")
+    ```
+    """
+    )
+
+    workspace = {}
+    to_files(chat, workspace)
+
+    assert workspace["all_output.txt"] == chat
+
+    expected_files = {
+        "file1.py": 'print("Hello, World!")\n',
+        "README.md": "\nThis is a sample program.\n\n`file1.py`\n",
+    }
+
+    for file_name, file_content in expected_files.items():
+        assert workspace[file_name] == file_content