Added Docstrings for Improved Code Documentation

promptshell/ai_terminal_assistant.py

@@ -13,6 +13,13 @@
 
 class AITerminalAssistant:
     def __init__(self, model_name: str, max_tokens: int = 8000, config: dict = None):
+        """Initializes the AITerminalAssistant with the given model name and configuration.
+
+        Args:
+            model_name (str): The name of the model to use for command execution.
+            max_tokens (int, optional): The maximum number of tokens for the model. Defaults to 8000.
+            config (dict, optional): Additional configuration settings. Defaults to None.
+        """
         self.username = getpass.getuser()
         self.home_folder = os.path.expanduser("~")
         self.current_directory = os.getcwd()
@@ -28,6 +35,7 @@ def __init__(self, model_name: str, max_tokens: int = 8000, config: dict = None)
         self.initialize_system_context()
 
     def initialize_system_context(self):
+        """Initializes the system context by gathering installed commands and system information."""
         path_dirs = os.environ.get('PATH', '').split(os.pathsep)
         installed_commands = []
         for dir in path_dirs:
@@ -105,7 +113,7 @@ def initialize_system_context(self):
         Solution: Confirm file existence with 'ls'
         Alternative: Use trash-cli instead of rm
         """
-        
+
         self.debugger.definition = f"""
         [ROLE] Shell Environment Debugger
         [TASK] Diagnose complex system issues
@@ -154,6 +162,14 @@ def initialize_system_context(self):
         """
 
     def execute_command_with_live_output(self, command: str) -> Tuple[str, str, int]:
+        """Executes a shell command and captures its live output.
+
+        Args:
+            command (str): The shell command to execute.
+
+        Returns:
+            Tuple[str, str, int]: A tuple containing the standard output, standard error, and exit code.
+        """
         interactive_commands = [
             'vim', 'vi', 'nano', 'emacs', 'ssh', 'telnet', 'top', 'htop',
             'man', 'less', 'more', 'mysql', 'psql', 'nmtui', 'crontab',
@@ -183,6 +199,14 @@ def execute_command_with_live_output(self, command: str) -> Tuple[str, str, int]
             return "", str(e), 1
 
     def execute_interactive_command(self, command: str) -> Tuple[str, str, int]:
+        """Executes an interactive shell command.
+
+        Args:
+            command (str): The interactive shell command to execute.
+
+        Returns:
+            Tuple[str, str, int]: A tuple containing the standard output, standard error, and exit code.
+        """        
         print(format_text('yellow') + "Executing interactive command..." + reset_format())
         try:
             proc = subprocess.Popen(
@@ -201,6 +225,14 @@ def execute_interactive_command(self, command: str) -> Tuple[str, str, int]:
             return "", str(e), 1
 
     def execute_command(self, user_input: str) -> str:
+        """Executes a command based on user input.
+
+        Args:
+            user_input (str): The command or question input by the user.
+
+        Returns:
+            str: The result of the command execution or an error message.
+        """
         try:
             self.current_directory = os.getcwd()
             if user_input.strip() == "":
@@ -260,6 +292,14 @@ def execute_command(self, user_input: str) -> str:
             return self.handle_error(str(e), user_input, command)
 
     def run_direct_command(self, command: str) -> str:
+        """Executes a direct shell command provided by the user.
+
+        Args:
+            command (str): The shell command to execute.
+
+        Returns:
+            str: The result of the command execution or an error message.
+        """
         try:
             formatted_command = format_text('cyan') + f"Direct Command: {command}" + reset_format()
             print(formatted_command)
@@ -287,6 +327,14 @@ def run_direct_command(self, command: str) -> str:
             return self.handle_error(str(e), command, command)
 
     def answer_question(self, question: str) -> str:
+        """Answers a user-provided question based on the current context.
+
+        Args:
+            question (str): The question to answer.
+
+        Returns:
+            str: The answer to the question.
+        """
         context = f"""
         Command History (last 10 commands):
         {', '.join(self.command_history)}
@@ -301,6 +349,14 @@ def answer_question(self, question: str) -> str:
         return format_text('cyan') + "Answer:\n" + answer + reset_format()
 
     def gather_additional_data(self, user_input: str) -> dict:
+        """Gathers additional data based on the user's input, such as clipboard content or file data.
+
+        Args:
+            user_input (str): The user's input to analyze for additional data needs.
+
+        Returns:
+            dict: A dictionary containing additional data, such as clipboard content or file content.
+        """
         additional_data = {}
         if "clipboard" in user_input.lower():
             clipboard_content = self.data_gatherer.get_clipboard_content()
@@ -318,6 +374,16 @@ def gather_additional_data(self, user_input: str) -> dict:
         return additional_data
 
     def debug_error(self, command: str, error_output: str, exit_code: int) -> str:
+        """Analyzes a failed command and provides debugging suggestions.
+
+        Args:
+            command (str): The command that failed.
+            error_output (str): The error output from the failed command.
+            exit_code (int): The exit code of the failed command.
+
+        Returns:
+            str: A debugging suggestion or alternative command.
+        """
         context = f"""
         Command History (last 10 commands):
         {', '.join(self.command_history)}
@@ -335,6 +401,16 @@ def debug_error(self, command: str, error_output: str, exit_code: int) -> str:
         return self.debugger(debug_input)
 
     def handle_error(self, error: str, user_input: str, command: str) -> str:
+        """Handles errors by analyzing the issue and suggesting a corrected command.
+
+        Args:
+            error (str): The error message.
+            user_input (str): The original user input.
+            command (str): The interpreted command that caused the error.
+
+        Returns:
+            str: The result of executing the suggested command or an error message.
+        """
         error_analysis = self.error_handler(f"""
         Error: {error}
         User Input: {user_input}

promptshell/data_gatherer.py

@@ -4,13 +4,26 @@
 class DataGatherer:
     @staticmethod
     def get_clipboard_content():
+        """Retrieves the current content of the clipboard.
+
+        Returns:
+            str: The content of the clipboard or an error message if unable to access it.
+        """
         try:
             return pyperclip.paste()
         except:
             return "Error: Unable to access clipboard"
 
     @staticmethod
     def get_file_content(file_path):
+        """Reads the content of a specified file.
+
+        Args:
+            file_path (str): The path to the file to read.
+
+        Returns:
+            str: The content of the file or an error message if reading fails.
+        """
         try:
             with open(file_path, 'r') as file:
                 return file.read()
@@ -19,6 +32,14 @@ def get_file_content(file_path):
 
     @staticmethod
     def execute_command(command):
+        """Executes a shell command and captures its output.
+
+        Args:
+            command (str): The shell command to execute.
+
+        Returns:
+            str: The standard output of the command or an error message if execution fails.
+        """
         try:
             result = subprocess.run(command, capture_output=True, text=True, shell=True)
             return result.stdout if result.returncode == 0 else f"Error: {result.stderr}"

promptshell/format_utils.py

@@ -2,6 +2,17 @@
 import platform
 
 def format_text(fg, bg=None, inverted=False, bold=False):
+    """Formats text for terminal output with specified foreground and background colors.
+
+    Args:
+        fg (str): The foreground color (e.g., 'red', 'green').
+        bg (str, optional): The background color (e.g., 'black', 'white'). Defaults to None.
+        inverted (bool, optional): If True, inverts the foreground and background colors. Defaults to False.
+        bold (bool, optional): If True, makes the text bold. Defaults to False.
+
+    Returns:
+        str: The formatted text string with ANSI escape codes.
+    """
     reset = "\033[0m"
     result = reset
     if bold:
@@ -18,22 +29,40 @@ def format_text(fg, bg=None, inverted=False, bold=False):
     return result
 
 def reset_format():
+    """Resets the text formatting to default.
+
+    Returns:
+        str: The ANSI escape code to reset formatting.
+    """
     return "\033[0m"
 
 def get_terminal_size():
+    """Retrieves the current size of the terminal window.
+
+    Returns:
+        tuple: A tuple containing the number of columns and rows in the terminal.
+    """
     try:
         columns, rows = os.get_terminal_size(0)
     except OSError:
         columns, rows = os.get_terminal_size(1)
     return columns, rows
 
 def get_current_os():
-    """Detect and normalize current operating system"""
+    """Detects and normalizes the current operating system.
+
+    Returns:
+        str: The name of the current operating system ('windows', 'macos', or 'linux').
+    """
     system = platform.system().lower()
     return 'windows' if system == 'windows' else 'macos' if system == 'darwin' else 'linux'
 
 def get_os_specific_examples():
-    """Return OS-appropriate command examples"""
+    """Returns OS-appropriate command examples based on the current operating system.
+
+    Returns:
+        list: A list of command examples specific to the current OS.
+    """
     current_os = get_current_os()
     examples = {
         'windows': [

promptshell/main.py

@@ -7,6 +7,20 @@
 from .setup import setup_wizard, load_config, get_active_model
 
 def main():
+    """The main entry point for the AI-Powered Terminal Assistant.
+
+    This function initializes the assistant, loads the configuration, and starts the interactive loop
+    for processing user input. It handles commands, questions, and configuration updates.
+
+    Behavior:
+        - Loads configuration or runs a setup wizard if no configuration exists.
+        - Enables ANSI support and sets up readline for command-line enhancements.
+        - Processes user input for commands, questions, or special options like '--help' or '--config'.
+        - Provides a clean exit on 'quit' or 'Ctrl + c'.
+
+    Returns:
+        None
+    """
     config = load_config()
     if not config:
         print("First-time setup required!")