edsl 0.1.47__py3-none-any.whl → 0.1.48__py3-none-any.whl

edsl/scenarios/directory_scanner.py

@@ -1,15 +1,56 @@
-# directory_scanner.py
+"""
+The DirectoryScanner module provides functionality for finding and processing files in directories.
+
+This module implements the DirectoryScanner class, which is designed to scan directories
+for files matching specific criteria and process them using a factory function. It supports
+recursive scanning, filtering by file extensions, and both eager and lazy iteration over
+the matching files.
+"""
+
 from dataclasses import dataclass
 from typing import Optional, List, Iterator, TypeVar, Generic, Callable, Any
 import os
+from pathlib import Path
 
+# Generic type variable for the factory function's return type
 T = TypeVar("T")
 
 
 @dataclass
 class DirectoryScanner:
     """
-    Scanner for finding files in a directory based on various criteria.
+    A utility class for finding and processing files in directories.
+    
+    DirectoryScanner provides methods to scan directories for files that match specific
+    criteria, such as file extensions. It can process matching files using a factory
+    function that converts file paths to objects of a specified type.
+    
+    The scanner supports both eager (scan) and lazy (iter_scan) iteration, recursive
+    directory traversal, and flexible filtering based on file extensions.
+    
+    Attributes:
+        directory_path: The path to the directory to scan.
+        
+    Examples:
+        >>> import tempfile
+        >>> import os
+        >>> # Create a temporary directory with some files
+        >>> with tempfile.TemporaryDirectory() as tmpdir:
+        ...     # Create a few files with different extensions
+        ...     _ = open(os.path.join(tmpdir, "file1.txt"), "w").write("content")
+        ...     _ = open(os.path.join(tmpdir, "file2.txt"), "w").write("content")
+        ...     _ = open(os.path.join(tmpdir, "image.jpg"), "w").write("content")
+        ...     # Create a scanner and find all text files
+        ...     scanner = DirectoryScanner(tmpdir)
+        ...     txt_files = scanner.scan(lambda path: path, suffix_allow_list=["txt"])
+        ...     len(txt_files)
+        ...     # Use a factory to process files
+        ...     def get_filename(path):
+        ...         return os.path.basename(path)
+        ...     filenames = scanner.scan(get_filename)
+        ...     sorted(filenames)
+        2
+        ['file1.txt', 'file2.txt', 'image.jpg']
     """
 
     directory_path: str
@@ -24,15 +65,56 @@ class DirectoryScanner:
         include_no_extension: bool = True,
     ) -> List[T]:
         """
-        Eagerly scan directory and return list of objects created by factory.
-
+        Eagerly scan directory and return a list of objects created by the factory function.
+        
+        This method performs a scan of the directory, filtering files based on the provided
+        criteria, and applies the factory function to each matching file path. It returns
+        a complete list of processed results.
+        
         Args:
-            factory: Callable that creates objects from file paths
-            recursive: If True, recursively traverse subdirectories
-            suffix_allow_list: List of allowed file extensions (without dots)
-            suffix_exclude_list: List of excluded file extensions (takes precedence over allow list)
-            example_suffix: If provided, only include files with this example suffix
-            include_no_extension: Whether to include files without extensions
+            factory: A callable that takes a file path string and returns an object of type T.
+                    This is applied to each matching file path.
+            recursive: If True, traverses subdirectories recursively. If False, only scans
+                      the top-level directory.
+            suffix_allow_list: A list of file extensions (without dots) to include.
+                              If provided, only files with these extensions are included.
+            suffix_exclude_list: A list of file extensions to exclude. This takes precedence
+                                over suffix_allow_list.
+            example_suffix: If provided, only include files ending with this exact suffix.
+                           This checks the entire filename, not just the extension.
+            include_no_extension: Whether to include files without extensions. Defaults to True.
+            
+        Returns:
+            A list of objects created by applying the factory function to each matching file path.
+            
+        Examples:
+            >>> import tempfile
+            >>> import os
+            >>> with tempfile.TemporaryDirectory() as tmpdir:
+            ...     # Create test files
+            ...     _ = open(os.path.join(tmpdir, "doc1.txt"), "w").write("content")
+            ...     _ = open(os.path.join(tmpdir, "doc2.md"), "w").write("content")
+            ...     os.mkdir(os.path.join(tmpdir, "subdir"))
+            ...     _ = open(os.path.join(tmpdir, "subdir", "doc3.txt"), "w").write("content")
+            ...     # Scan for text files only
+            ...     scanner = DirectoryScanner(tmpdir)
+            ...     paths = scanner.scan(lambda p: p, suffix_allow_list=["txt"])
+            ...     len(paths)
+            ...     # Recursive scan for all files
+            ...     all_paths = scanner.scan(lambda p: p, recursive=True)
+            ...     len(all_paths)
+            ...     # Exclude specific extensions
+            ...     no_md = scanner.scan(lambda p: p, recursive=True, suffix_exclude_list=["md"])
+            ...     len(no_md)
+            1
+            3
+            2
+            
+        Notes:
+            - This method is eager and collects all results into memory. For large directories,
+              consider using iter_scan instead.
+            - The filtering logic applies filters in this order: exclude list, example suffix,
+              allow list, and no extension.
         """
         return list(
             self.iter_scan(
@@ -55,12 +137,64 @@ class DirectoryScanner:
         include_no_extension: bool = True,
     ) -> Iterator[T]:
         """
-        Lazily scan directory and yield objects created by factory.
+        Lazily scan directory and yield objects created by the factory function.
+        
+        This method performs a lazy scan of the directory, filtering files based on the provided
+        criteria, and applies the factory function to each matching file path. It yields
+        results one by one, allowing for memory-efficient processing of large directories.
+        
+        Args:
+            factory: A callable that takes a file path string and returns an object of type T.
+                    This is applied to each matching file path.
+            recursive: If True, traverses subdirectories recursively. If False, only scans
+                      the top-level directory.
+            suffix_allow_list: A list of file extensions (without dots) to include.
+                              If provided, only files with these extensions are included.
+            suffix_exclude_list: A list of file extensions to exclude. This takes precedence
+                                over suffix_allow_list.
+            example_suffix: If provided, only include files ending with this exact suffix.
+                           This checks the entire filename, not just the extension.
+            include_no_extension: Whether to include files without extensions. Defaults to True.
+            
+        Yields:
+            Objects created by applying the factory function to each matching file path,
+            yielded one at a time.
+            
+        Examples:
+            >>> import tempfile
+            >>> import os
+            >>> with tempfile.TemporaryDirectory() as tmpdir:
+            ...     # Create test files
+            ...     _ = open(os.path.join(tmpdir, "doc1.txt"), "w").write("content")
+            ...     _ = open(os.path.join(tmpdir, "doc2.md"), "w").write("content")
+            ...     # Process files lazily
+            ...     scanner = DirectoryScanner(tmpdir)
+            ...     for path in scanner.iter_scan(lambda p: p):
+            ...         # Process each file path without loading all into memory
+            ...         file_exists = os.path.exists(path)
+            ...         assert file_exists
+            
+        Notes:
+            - This method is lazy and yields results as they are processed, making it
+              suitable for memory-efficient processing of large directories.
+            - The filtering logic is identical to the scan method.
         """
 
         def should_include_file(filepath: str) -> bool:
+            """
+            Determine if a file should be included based on filtering criteria.
+            
+            This helper function applies all the filtering rules to determine
+            if a given file path should be included in the results.
+            
+            Args:
+                filepath: The path to the file to check.
+                
+            Returns:
+                True if the file should be included, False otherwise.
+            """
             _, ext = os.path.splitext(filepath)
-            ext = ext[1:] if ext else ""
+            ext = ext[1:] if ext else ""  # Remove leading dot from extension
 
             # Handle no extension case
             if not ext:
@@ -80,7 +214,16 @@ class DirectoryScanner:
 
             return True
 
-        def iter_files():
+        def iter_files() -> Iterator[str]:
+            """
+            Generate paths to all files in the directory, optionally recursively.
+            
+            This helper function yields file paths from the directory, handling
+            the recursive option appropriately.
+            
+            Yields:
+                Paths to files in the directory.
+            """
             if recursive:
                 for root, _, files in os.walk(self.directory_path):
                     for file in files:

edsl/scenarios/document_chunker.py

@@ -0,0 +1,186 @@
+"""
+The DocumentChunker module provides functionality for splitting text into manageable chunks.
+
+This module implements the DocumentChunker class, which is responsible for chunking
+text content in Scenarios based on word or line counts. This is particularly useful
+when working with large text documents that need to be processed in smaller pieces,
+such as for summarization, analysis, or when dealing with models that have token
+limits.
+"""
+
+from __future__ import annotations
+from typing import Optional, Generator, TYPE_CHECKING, List, Union
+import copy
+import hashlib
+
+from .scenario import Scenario
+from .scenario_list import ScenarioList
+
+
+class DocumentChunker:
+    """
+    A utility class for splitting text in a Scenario into manageable chunks.
+    
+    DocumentChunker provides methods to split text content from a Scenario field into
+    smaller chunks based on either word count or line count. It's primarily used by the
+    Scenario.chunk() method but can also be used directly for more control over the
+    chunking process.
+    
+    Attributes:
+        scenario: The Scenario object containing the text to be chunked.
+    """
+    
+    def __init__(self, scenario: "Scenario"):
+        """
+        Initialize a DocumentChunker for a specific Scenario.
+        
+        Args:
+            scenario: The Scenario object containing the text field to be chunked.
+        """
+        self.scenario = scenario
+
+    @staticmethod
+    def _line_chunks(text: str, num_lines: int) -> Generator[str, None, None]:
+        """
+        Split text into chunks based on a specified number of lines per chunk.
+        
+        This method divides a text string into chunks, where each chunk contains
+        at most the specified number of lines. It processes the text by splitting
+        on newline characters and then groups the lines into chunks.
+        
+        Args:
+            text: The text string to split into chunks.
+            num_lines: The maximum number of lines to include in each chunk.
+            
+        Yields:
+            String chunks containing at most num_lines lines each.
+            
+        Examples:
+            >>> list(DocumentChunker._line_chunks("This is a test.\\nThis is a test. This is a test.", 1))
+            ['This is a test.', 'This is a test. This is a test.']
+            
+            >>> list(DocumentChunker._line_chunks("Line 1\\nLine 2\\nLine 3\\nLine 4", 2))
+            ['Line 1\\nLine 2', 'Line 3\\nLine 4']
+        """
+        lines = text.split("\n")
+        for i in range(0, len(lines), num_lines):
+            chunk = "\n".join(lines[i : i + num_lines])
+            yield chunk
+
+    @staticmethod
+    def _word_chunks(text: str, num_words: int) -> Generator[str, None, None]:
+        """
+        Split text into chunks based on a specified number of words per chunk.
+        
+        This method divides a text string into chunks, where each chunk contains
+        at most the specified number of words. It processes the text by splitting
+        on whitespace and then groups the words into chunks.
+        
+        Args:
+            text: The text string to split into chunks.
+            num_words: The maximum number of words to include in each chunk.
+            
+        Yields:
+            String chunks containing at most num_words words each.
+            
+        Examples:
+            >>> list(DocumentChunker._word_chunks("This is a test.", 2))
+            ['This is', 'a test.']
+            
+            >>> list(DocumentChunker._word_chunks("One two three four five", 3))
+            ['One two three', 'four five']
+        """
+        words = text.split()
+        for i in range(0, len(words), num_words):
+            chunk = " ".join(words[i : i + num_words])
+            yield chunk
+
+    def chunk(
+        self,
+        field: str,
+        num_words: Optional[int] = None,
+        num_lines: Optional[int] = None,
+        include_original: bool = False,
+        hash_original: bool = False,
+    ) -> ScenarioList:
+        """
+        Split a text field in the Scenario into chunks and create a ScenarioList.
+        
+        This method takes a field containing text from the Scenario and divides it into
+        smaller chunks based on either word count or line count. For each chunk, it creates
+        a new Scenario with additional metadata about the chunk.
+        
+        Args:
+            field: The key name of the field in the Scenario to split.
+            num_words: The number of words to include in each chunk. Mutually exclusive
+                      with num_lines.
+            num_lines: The number of lines to include in each chunk. Mutually exclusive
+                      with num_words.
+            include_original: If True, includes the original complete text in each chunk
+                             with a "_original" suffix.
+            hash_original: If True and include_original is True, stores a hash of the
+                          original text instead of the full text.
+        
+        Returns:
+            A ScenarioList containing multiple Scenarios, each with a chunk of the
+            original text and metadata about the chunk.
+            
+        Raises:
+            ValueError: If neither num_words nor num_lines is specified, or if both are.
+            KeyError: If the specified field doesn't exist in the Scenario.
+            
+        Notes:
+            - Each chunk is assigned a sequential index in the '{field}_chunk' field
+            - Character and word counts for each chunk are included in '{field}_char_count'
+              and '{field}_word_count' fields
+            - When include_original is True, the original text is preserved in each chunk
+              in the '{field}_original' field
+            - The hash_original option is useful to save space while maintaining traceability
+        """
+        # Check if field exists in the scenario
+        if field not in self.scenario:
+            raise KeyError(f"Field '{field}' not found in the scenario")
+
+        # Validate parameters
+        if num_words is None and num_lines is None:
+            raise ValueError("You must specify either num_words or num_lines.")
+
+        if num_words is not None and num_lines is not None:
+            raise ValueError(
+                "You must specify either num_words or num_lines, but not both."
+            )
+            
+        # Get appropriate chunks based on the specified chunking method
+        if num_words is not None:
+            chunks = list(self._word_chunks(self.scenario[field], num_words))
+        else:  # num_lines is not None
+            chunks = list(self._line_chunks(self.scenario[field], num_lines))
+
+        # Create a new scenario for each chunk with metadata
+        scenarios = []
+        for i, chunk in enumerate(chunks):
+            new_scenario = copy.deepcopy(self.scenario)
+            new_scenario[field] = chunk
+            new_scenario[field + "_chunk"] = i
+            new_scenario[field + "_char_count"] = len(chunk)
+            new_scenario[field + "_word_count"] = len(chunk.split())
+            
+            # Include the original text if requested
+            if include_original:
+                if hash_original:
+                    # Use MD5 hash for brevity, not for cryptographic security
+                    new_scenario[field + "_original"] = hashlib.md5(
+                        self.scenario[field].encode()
+                    ).hexdigest()
+                else:
+                    new_scenario[field + "_original"] = self.scenario[field]
+                    
+            scenarios.append(new_scenario)
+            
+        return ScenarioList(scenarios)
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

edsl/scenarios/exceptions.py

@@ -0,0 +1,101 @@
+"""
+Exceptions module for the scenarios package.
+
+This module defines custom exception classes used throughout the scenarios module.
+These exceptions provide specific error information for different types of errors
+that can occur when working with Scenarios, ScenarioLists, and related components.
+"""
+
+import re
+from typing import List
+
+from ..base import BaseException
+
+
+class AgentListError(BaseException):
+    """
+    Exception raised for errors related to AgentList operations.
+    
+    This exception is raised when there are issues with creating, modifying,
+    or using an AgentList in conjunction with scenarios.
+    
+    Args:
+        message: A description of the error that occurred.
+    """
+    
+    def __init__(self, message: str):
+        """
+        Initialize the AgentListError with a message.
+        
+        Args:
+            message: A description of the error that occurred.
+        """
+        super().__init__(message)
+
+
+class ScenarioError(BaseException):
+    """
+    Exception raised for errors related to Scenario operations.
+    
+    This exception is raised when there are issues with creating, modifying,
+    or using Scenarios. It automatically includes a link to the documentation
+    in the error message and makes URLs clickable in terminal output.
+    
+    Args:
+        message: A description of the error that occurred.
+    """
+    
+    documentation = "https://docs.expectedparrot.com/en/latest/scenarios.html#module-edsl.scenarios.Scenario"
+
+    def __init__(self, message: str):
+        """
+        Initialize the ScenarioError with a message and add documentation link.
+        
+        Args:
+            message: A description of the error that occurred.
+        """
+        self.message = message + "\n" + "Documentation: " + self.documentation
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the error with clickable URLs.
+        
+        This method makes any URLs in the error message clickable when displayed
+        in terminal environments that support ANSI escape sequences.
+        
+        Returns:
+            The error message with clickable URLs.
+        """
+        return self.make_urls_clickable(self.message)
+
+    @staticmethod
+    def make_urls_clickable(text: str) -> str:
+        """
+        Convert URLs in text to clickable links in terminal output.
+        
+        This method finds all URLs in the given text and wraps them in ANSI
+        escape sequences that make them clickable in supporting terminals.
+        
+        Args:
+            text: The text containing URLs to make clickable.
+            
+        Returns:
+            The text with URLs converted to clickable links.
+            
+        Example:
+            >>> error = ScenarioError("See docs at https://example.com")
+            >>> s = str(error)  # Returns the message with clickable link
+            ...
+        """
+        url_pattern = r"https?://[^\s]+"
+        urls = re.findall(url_pattern, text)
+        for url in urls:
+            clickable_url = f"\033]8;;{url}\007{url}\033]8;;\007"
+            text = text.replace(url, clickable_url)
+        return text
+
+
+if __name__ == "__main__":
+    import doctest
+    doctest.testmod(optionflags=doctest.ELLIPSIS)

edsl/scenarios/file_methods.py

@@ -2,8 +2,7 @@ from typing import Optional, Dict, Type
 from abc import ABC, abstractmethod
 import importlib.metadata
 import importlib.util
-
-from edsl.utilities.is_notebook import is_notebook
+from ..utilities import is_notebook
 
 
 class FileMethods(ABC):
@@ -30,7 +29,7 @@ class FileMethods(ABC):
     def load_plugins(cls):
         """Load all file handler plugins including built-ins and external plugins."""
 
-        from edsl.scenarios import handlers
+        from . import handlers  # noqa: F401 - import needed for handler registration
 
         # Then load any external plugins
         try: