Enhance file handling and testing in Coder class 📁🔍

Resolves #57

This commit introduces several improvements to the `Coder` class in `aicodebot/coder.py`. We've added functionality to better handle different file types, including binary files and files with unknown types. This includes the addition of the `get_file_info` and `is_binary_file` methods. We've also updated the handling of new files in the `git_diff_context` method to account for binary files.

In addition, we've added a new binary file `assets/robot.png` to the repository.

Finally, we've expanded our test coverage in `tests/test_coder.py` to include tests for the new `get_file_info` method. This ensures our file handling remains robust and reliable. 🧪👍

Author: TechNickAI

aicodebot/coder.py

@@ -4,7 +4,8 @@
 from openai.api_resources import engine
 from pathlib import Path
 from prompt_toolkit.completion import Completer, Completion
-import fnmatch, functools, openai, os, re, subprocess, tiktoken
+from pygments.lexers import ClassNotFound, get_lexer_for_mimetype, guess_lexer_for_filename
+import fnmatch, functools, mimetypes, openai, os, re, subprocess, tiktoken
 
 DEFAULT_MAX_TOKENS = 512
 PRECISE_TEMPERATURE = 0.05
@@ -17,6 +18,8 @@ class Coder:
     git, and the local file system.
     """
 
+    UNKNOWN_FILE_TYPE = "unknown"
+
     @staticmethod
     def clone_repo(repo_url, repo_dir):
         """Clone a git repository to a directory."""
@@ -79,6 +82,31 @@ def generate_directory_structure(cls, path, ignore_patterns=None, use_gitignore=
 
         return structure
 
+    @classmethod
+    def get_file_info(cls, file_path):
+        # Use the mimetypes module to guess the MIME type
+        mime_type = mimetypes.guess_type(file_path)[0]
+
+        # Use the is_binary_file function to check if the file is binary
+        is_binary = cls.is_binary_file(file_path)
+
+        try:
+            # Try to get the lexer for the MIME type
+            lexer = get_lexer_for_mimetype(mime_type)
+        except ClassNotFound:
+            try:
+                # If that fails, try to guess the lexer based on the file name
+                lexer = guess_lexer_for_filename(file_path, "")
+            except ClassNotFound:
+                # If that also fails, set the file type to UNKNOWN_FILE_TYPE
+                file_type = cls.UNKNOWN_FILE_TYPE
+            else:
+                file_type = lexer.name
+        else:
+            file_type = lexer.name
+
+        return is_binary, file_type
+
     @staticmethod
     @functools.lru_cache
     def get_openai_supported_engines():
@@ -230,9 +258,13 @@ def git_diff_context(commit=None, files=None):
                 if status_code == "A":
                     # If the file is new, include the entire file content
                     file_name = status_parts[1]
-                    contents = Path(file_name).read_text()
-                    diffs.append(f"## New file added: {file_name}")
-                    diffs.append(contents)
+                    if Coder.is_binary_file(file_name):
+                        # Don't include the diff for binary files
+                        diffs.append(f"## New binary file added: {file_name}")
+                    else:
+                        diffs.append(f"## New file added: {file_name}")
+                        contents = Path(file_name).read_text()
+                        diffs.append(contents)
                 elif status_code == "R":
                     # If the file is renamed, get the diff and note the old and new names
                     old_file_name, new_file_name = status_parts[1], status_parts[2]
@@ -258,6 +290,18 @@ def git_staged_files():
     def git_unstaged_files():
         return exec_and_get_output(["git", "diff", "HEAD", "--name-only"]).splitlines()
 
+    @staticmethod
+    def is_binary_file(file_path):
+        chunksize = 4000
+        with Path(file_path).open("rb") as file:
+            while True:
+                chunk = file.read(chunksize)
+                if b"\0" in chunk:  # Null byte
+                    return True
+                if len(chunk) < chunksize:
+                    break  # End of file
+        return False
+
     @staticmethod
     def parse_github_url(repo_url):
         """

assets/robot.png

@@ -44,6 +44,34 @@ def test_generate_directory_structure(tmp_path):
     assert len(file_list) > 10
 
 
+def test_get_file_info():
+    # Test with a text file
+    is_binary, file_type = Coder.get_file_info("tests/test_coder.py")
+    assert is_binary is False
+    assert file_type == "Python"
+
+    # Test with a binary file
+    is_binary, file_type = Coder.get_file_info("assets/robot.png")
+    assert is_binary is True
+    assert file_type == Coder.UNKNOWN_FILE_TYPE
+
+    is_binary, file_type = Coder.get_file_info("LICENSE")
+    assert is_binary is False
+    assert file_type == Coder.UNKNOWN_FILE_TYPE
+
+    is_binary, file_type = Coder.get_file_info("README.md")
+    assert is_binary is False
+    assert file_type == "Markdown"
+
+    is_binary, file_type = Coder.get_file_info("pyproject.toml")
+    assert is_binary is False
+    assert file_type == "TOML"
+
+    # Test with a non-existent file
+    with pytest.raises(FileNotFoundError):
+        Coder.get_file_info("non_existent_file.txt")
+
+
 def test_get_token_length():
     text = ""
     assert Coder.get_token_length(text) == 0