capture.py

# --------------------------------------------------------------------------------------
# Capture Class - Output Redirection System
# --------------------------------------------------------------------------------------
# PURPOSE: Provides functionality to capture print() output into memory
#          This allows saving console output to files without displaying it on screen
#
# USE CASE: When generating reports, we want to:
#           1. Run all report functions
#           2. Capture their output (instead of printing to console)
#           3. Write the captured output to a file
# --------------------------------------------------------------------------------------

class Capture:
    def __init__(self):
        """Initialize the capture system with default print function."""
        self.print_func = print  # Default: use normal Python print
        
    # ----------------------------------------------------------------------------------
    # PRINT CAPTURE SYSTEM
    # ----------------------------------------------------------------------------------
    # These methods redirect print() output to memory for later file writing
    
    def _format_args(self, *args, sep=" ", end="\n"):
        """
        Format print arguments into a single string.
        
        Mimics Python's built-in print() behavior:
            - Joins arguments with separator (default: space)
            - Adds end character (default: newline)
        
        PARAMETERS:
            *args: Values to print (like regular print)
            sep: Separator between values
            end: String appended after last value
        
        RETURNS: Formatted string ready for storage
        """
        return sep.join(str(a) for a in args) + ("" if end is None else end)
    
    def _capture_print(self, *args, sep=" ", end="\n", **kwargs):
        """
        Redirect print() output into memory instead of console.
        
        This is used as a replacement for Python's print() function when we want
        to save output to a file instead of displaying it.
        
        HOW IT WORKS:
            1. Format the arguments like normal print()
            2. Store the result in collected_print_lines list
            3. Strip trailing newline (we'll add them back when writing to file)
        """
        line = self._format_args(*args, sep=sep, end=end)
        if line.endswith("\n"):
            line = line[:-1]  # Remove newline for cleaner storage
        self.collected_print_lines.append(line)
    
    def set_print_func(self, func):
        """
        Switch between normal print and captured print.
        
        PARAMETERS:
            func: Callable to use as print function
                  - Pass 'print' for normal console output
                  - Pass '_capture_print' for memory capture
        """
        if not callable(func):
            raise ValueError("print function must be callable")
        self.print_func = func
        self.capture_enabled = False
    
    def use_capture(self, enable=True, clear_existing=False):
        """
        Enable or disable output capturing.
        
        WHEN ENABLED:
            - All calls to self.print_func() will store output in memory
            - Nothing appears on the console
        
        WHEN DISABLED:
            - self.print_func() works like normal print()
            - Output appears on console
        
        PARAMETERS:
            enable: True = capture output, False = print normally
            clear_existing: If True, clears previously captured lines
        """
        if enable:
            self.print_func = self._capture_print
            self.capture_enabled = True
            if clear_existing:
                self.collected_print_lines = []
        else:
            self.print_func = print
            self.capture_enabled = False
    
    def get_captured(self):
        """
        Get all captured print lines.
        
        RETURNS: Copy of the list containing all captured output lines
        """
        return list(self.collected_print_lines)
    
    def clear_captured(self):
        """
        Clear the buffer of captured print lines.
        
        Useful for starting fresh capture or freeing memory after saving to file.
        """
        self.collected_print_lines = []