Add directory structure generator and refactor tests

Alright, listen up. We've got some new shit in town. We've added a badass function to generate a directory structure, and it's got all the bells and whistles. It respects .gitignore, it can ignore patterns, and it's got a pretty print for your viewing pleasure.

We've also been cleaning up the house. Moved the create_and_write_file function to conftest.py, cause that's where it belongs. It's a utility function, and it's gonna be used in multiple test files.

Speaking of tests, we've added a whole bunch of them for the new directory structure generator. We're making sure this function is bulletproof.

And lastly, we've done some minor text changes in prompts.py and conftest.py. Nothing to lose your shit over.

So there you have it. New features, better tests, cleaner code. That's how we roll.

Author: TechNickAI

aicodebot/helpers.py

@@ -1,6 +1,6 @@
 from loguru import logger
 from pathlib import Path
-import os, subprocess, sys, tiktoken, yaml
+import fnmatch, os, subprocess, sys, tiktoken, yaml
 
 # ---------------------------------------------------------------------------- #
 #                    Global logging configuration for loguru                   #
@@ -20,6 +20,31 @@
 # ---------------------------------------------------------------------------- #
 
 
+def generate_directory_structure(path, ignore_patterns=None, use_gitignore=True, indent=0):
+    """Generate a text representation of the directory structure of a path."""
+    ignore_patterns = ignore_patterns.copy() if ignore_patterns else []
+
+    base_path = Path(path)
+
+    if use_gitignore:
+        # Note: .gitignore files can exist in sub directories as well, such as * in __pycache__ directories
+        gitignore_file = base_path / ".gitignore"
+        if gitignore_file.exists():
+            with gitignore_file.open() as f:
+                ignore_patterns.extend(line.strip() for line in f if line.strip() and not line.startswith("#"))
+
+    structure = ""
+    if base_path.is_dir():
+        if not any(fnmatch.fnmatch(base_path.name, pattern) for pattern in ignore_patterns):
+            structure += "  " * indent + f"- [Directory] {base_path.name}\n"
+            for item in base_path.iterdir():
+                structure += generate_directory_structure(item, ignore_patterns, use_gitignore, indent + 1)
+    elif not any(fnmatch.fnmatch(base_path.name, pattern) for pattern in ignore_patterns):
+        structure += "  " * indent + f"- [File] {base_path.name}\n"
+
+    return structure
+
+
 def get_llm_model(token_size=0):
     # https://platform.openai.com/docs/models/gpt-3-5
     # We want to use GPT-4, if it is available for this OPENAI_API_KEY, otherwise GPT-3.5

aicodebot/prompts.py

@@ -112,9 +112,7 @@ def generate_files_context(files):
 # ---------------------------------------------------------------------------- #
 
 ALIGNMENT_TEMPLATE = (
-    """
-    You're an advocate for aligned AI.
-    """
+    """You're an advocate for aligned AI."""
     + get_personality_prompt()
     + """
     You don't subscribe to the idea that AI is a black box or follow the

tests/conftest.py

@@ -10,6 +10,11 @@ def cli_runner():
     return CliRunner()
 
 
+def create_and_write_file(filename, text):
+    with Path(filename).open("w") as f:
+        f.write(text)
+
+
 @pytest.fixture
 def temp_git_repo(tmp_path):
     # Create a temporary git repository that can be used for testing

tests/test_helpers.py

@@ -1,8 +1,46 @@
-from aicodebot.helpers import exec_and_get_output, get_token_length, git_diff_context
-from pathlib import Path
+from aicodebot.helpers import exec_and_get_output, generate_directory_structure, get_token_length, git_diff_context
+from tests.conftest import create_and_write_file
 import os, pytest
 
 
+def test_generate_directory_structure(tmp_path):
+    # Create a file, a hidden file, another file, a .gitignore file, and a subdirectory in the temporary directory
+    create_and_write_file(tmp_path / "file.txt", "This is a test file")
+    create_and_write_file(tmp_path / ".hidden_file", "This is a hidden test file")
+    create_and_write_file(tmp_path / "test_file", "This is another test file")
+    create_and_write_file(tmp_path / ".gitignore", "*.txt\n")
+    sub_dir = tmp_path / "sub_dir"
+    sub_dir.mkdir()
+    create_and_write_file(sub_dir / "sub_file", "This is a test file in a subdirectory")
+    create_and_write_file(sub_dir / ".gitignore", "sub_file")
+
+    # Call the function with the temporary directory and an ignore pattern
+    directory_structure = generate_directory_structure(tmp_path, ignore_patterns=["*.txt"])
+
+    # Check that the returned string is not empty
+    assert directory_structure
+
+    # Check that the returned string contains the name of the created subdirectory
+    assert "- [Directory] sub_dir" in directory_structure
+
+    # Check that the returned string does not contain the names of the ignored files
+    assert "- [File] file.txt" not in directory_structure
+
+    # Check that the returned string contains the name of the hidden file and the other file
+    assert "- [File] .hidden_file" in directory_structure
+    assert "- [File] test_file" in directory_structure
+
+    # Check that the function respects .gitignore
+    directory_structure_gitignore = generate_directory_structure(tmp_path)
+    assert "- [File] file.txt" not in directory_structure_gitignore
+    assert "- [File] sub_file" not in directory_structure_gitignore
+
+    # Check that the function works correctly when use_gitignore is False
+    directory_structure_no_gitignore = generate_directory_structure(tmp_path, use_gitignore=False)
+    assert "- [File] file.txt" in directory_structure_no_gitignore
+    assert "- [File] sub_file" in directory_structure_no_gitignore
+
+
 def test_get_token_length():
     text = ""
     assert get_token_length(text) == 0
@@ -11,11 +49,6 @@ def test_get_token_length():
     assert get_token_length(text) == 14
 
 
-def create_and_write_file(filename, text):
-    with Path(filename).open("w") as f:
-        f.write(text)
-
-
 def test_git_diff_context(temp_git_repo):
     os.chdir(temp_git_repo.working_dir)
 

tests/test_workflow.py

@@ -1,15 +1,14 @@
 from aicodebot.cli import cli
 from git import Repo
-from pathlib import Path
+from tests.conftest import create_and_write_file
 import os, pytest
 
 
 @pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping live tests without an API key.")
 def test_commit_command(cli_runner, temp_git_repo):
     with cli_runner.isolated_filesystem():
         os.chdir(temp_git_repo.working_dir)  # change to the temporary repo directory
-        with Path.open("test.txt", "a") as f:
-            f.write("Adding a new line.")
+        create_and_write_file("test.txt", "This is a test file.")
         result = cli_runner.invoke(cli, ["commit", "-y"])
         assert result.exit_code == 0
         assert "The following files will be committed:\ntest.txt" in result.output
@@ -26,8 +25,7 @@ def test_review(cli_runner, temp_git_repo):
         os.chdir(temp_git_repo.working_dir)  # change to the temporary repo directory
 
         # Add a new file
-        with Path.open("test.txt", "w") as f:
-            f.write("Adding a new line.")
+        create_and_write_file("test.txt", "Adding a new line.")
 
         repo = Repo(temp_git_repo.working_dir)
         # Stage the new file