edsl 0.1.39__py3-none-any.whl → 0.1.39.dev1__py3-none-any.whl

edsl/scenarios/Scenario.py

@@ -2,65 +2,54 @@
 
 from __future__ import annotations
 import copy
+import hashlib
 import os
 import json
 from collections import UserDict
-from typing import Union, List, Optional, TYPE_CHECKING, Collection
+from typing import Union, List, Optional, Generator
 from uuid import uuid4
 
 from edsl.Base import Base
 from edsl.scenarios.ScenarioHtmlMixin import ScenarioHtmlMixin
-from edsl.utilities.remove_edsl_version import remove_edsl_version
+from edsl.utilities.decorators import add_edsl_version, remove_edsl_version
 from edsl.exceptions.scenarios import ScenarioError
 
-if TYPE_CHECKING:
-    from edsl.scenarios.ScenarioList import ScenarioList
-    from edsl.results.Dataset import Dataset
-
 
 class DisplayJSON:
-    """Display a dictionary as JSON."""
-
-    def __init__(self, input_dict: dict):
-        self.text = json.dumps(input_dict, indent=4)
+    def __init__(self, dict):
+        self.text = json.dumps(dict, indent=4)
 
     def __repr__(self):
         return self.text
 
 
 class DisplayYAML:
-    """Display a dictionary as YAML."""
-
-    def __init__(self, input_dict: dict):
+    def __init__(self, dict):
         import yaml
 
-        self.text = yaml.dump(input_dict)
+        self.text = yaml.dump(dict)
 
     def __repr__(self):
         return self.text
 
 
 class Scenario(Base, UserDict, ScenarioHtmlMixin):
-    """A Scenario is a dictionary of keys/values that can be used to parameterize questions."""
+    """A Scenario is a dictionary of keys/values.
+
+    They can be used parameterize EDSL questions."""
 
     __documentation__ = "https://docs.expectedparrot.com/en/latest/scenarios.html"
 
-    def __init__(self, data: Optional[dict] = None, name: Optional[str] = None):
+    def __init__(self, data: Union[dict, None] = None, name: str = None):
         """Initialize a new Scenario.
 
-        :param data: A dictionary of keys/values for parameterizing questions.
-        :param name: The name of the scenario.
-        """
+        # :param data: A dictionary of keys/values for parameterizing questions.
+        #"""
         if not isinstance(data, dict) and data is not None:
-            try:
-                data = dict(data)
-            except Exception as e:
-                raise ScenarioError(
-                    f"You must pass in a dictionary to initialize a Scenario. You passed in {data}",
-                    "Exception message:" + str(e),
-                )
+            raise EDSLScenarioError(
+                "You must pass in a dictionary to initialize a Scenario."
+            )
 
-        super().__init__()
         self.data = data if data is not None else {}
         self.name = name
 
@@ -70,6 +59,7 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         :param n: The number of times to replicate the scenario.
 
         Example:
+
         >>> s = Scenario({"food": "wood chips"})
         >>> s.replicate(2)
         ScenarioList([Scenario({'food': 'wood chips'}), Scenario({'food': 'wood chips'})])
@@ -92,13 +82,13 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
                     return True
         return False
 
-    def _convert_jinja_braces(
-        self, replacement_left: str = "<<", replacement_right: str = ">>"
+    def convert_jinja_braces(
+        self, replacement_left="<<", replacement_right=">>"
     ) -> Scenario:
         """Convert Jinja braces to some other character.
 
         >>> s = Scenario({"food": "I love {{wood chips}}"})
-        >>> s._convert_jinja_braces()
+        >>> s.convert_jinja_braces()
         Scenario({'food': 'I love <<wood chips>>'})
 
         """
@@ -112,7 +102,7 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
                 new_scenario[key] = value
         return new_scenario
 
-    def __add__(self, other_scenario: Scenario) -> Scenario:
+    def __add__(self, other_scenario: "Scenario") -> "Scenario":
         """Combine two scenarios by taking the union of their keys
 
         If the other scenario is None, then just return self.
@@ -137,14 +127,11 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
             return s
 
     def rename(
-        self,
-        old_name_or_replacement_dict: Union[str, dict[str, str]],
-        new_name: Optional[str] = None,
-    ) -> Scenario:
+        self, old_name_or_replacement_dict: dict, new_name: Optional[str] = None
+    ) -> "Scenario":
         """Rename the keys of a scenario.
 
-        :param old_name_or_replacement_dict: A dictionary of old keys to new keys *OR* a string of the old key.
-        :param new_name: The new name of the key.
+        :param replacement_dict: A dictionary of old keys to new keys.
 
         Example:
 
@@ -169,26 +156,13 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
                 new_scenario[key] = value
         return new_scenario
 
-    def new_column_names(self, new_names: List[str]) -> Scenario:
-        """Rename the keys of a scenario.
-
-        >>> s = Scenario({"food": "wood chips"})
-        >>> s.new_column_names(["food_preference"])
-        Scenario({'food_preference': 'wood chips'})
-        """
-        try:
-            assert len(new_names) == len(self.keys())
-        except AssertionError:
-            print("The number of new names must match the number of keys.")
-
-        new_scenario = Scenario()
-        for new_names, value in zip(new_names, self.values()):
-            new_scenario[new_names] = value
-        return new_scenario
-
     def table(self, tablefmt: str = "grid") -> str:
-        """Display a scenario as a table."""
-        return self.to_dataset().table(tablefmt=tablefmt)
+        from edsl.results.Dataset import Dataset
+
+        keys = [key for key, value in self.items()]
+        values = [value for key, value in self.items()]
+        d = Dataset([{"key": keys}, {"value": values}])
+        return d.table(tablefmt=tablefmt)
 
     def json(self):
         return DisplayJSON(self.to_dict(add_edsl_version=False))
@@ -198,7 +172,7 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
 
         return DisplayYAML(self.to_dict(add_edsl_version=False))
 
-    def to_dict(self, add_edsl_version: bool = True) -> dict:
+    def to_dict(self, add_edsl_version=True) -> dict:
         """Convert a scenario to a dictionary.
 
         Example:
@@ -226,7 +200,8 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         return d
 
     def __hash__(self) -> int:
-        """Return a hash of the scenario.
+        """
+        Return a hash of the scenario.
 
         Example:
 
@@ -238,23 +213,44 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
 
         return dict_hash(self.to_dict(add_edsl_version=False))
 
+    def print(self):
+        from rich import print_json
+        import json
+
+        print_json(json.dumps(self.to_dict()))
+
     def __repr__(self):
         return "Scenario(" + repr(self.data) + ")"
 
     def to_dataset(self) -> "Dataset":
-        """Convert a scenario to a dataset.
-
-        >>> s = Scenario({"food": "wood chips"})
-        >>> s.to_dataset()
-        Dataset([{'key': ['food']}, {'value': ['wood chips']}])
-        """
+        # d = Dataset([{'a.b':[1,2,3,4]}])
         from edsl.results.Dataset import Dataset
 
-        keys = list(self.keys())
-        values = list(self.values())
+        keys = [key for key, value in self.items()]
+        values = [value for key, value in self.items()]
         return Dataset([{"key": keys}, {"value": values}])
 
-    def select(self, list_of_keys: Collection[str]) -> "Scenario":
+    def _repr_html_(self):
+        from tabulate import tabulate
+        import reprlib
+
+        d = self.to_dict(add_edsl_version=False)
+        # return self.to_dataset()
+        r = reprlib.Repr()
+        r.maxstring = 70
+
+        data = [[k, r.repr(v)] for k, v in d.items()]
+        from tabulate import tabulate
+
+        if hasattr(self, "__documentation__"):
+            footer = f"<a href='{self.__documentation__}'>(docs)</a></p>"
+        else:
+            footer = ""
+
+        table = str(tabulate(data, headers=["keys", "values"], tablefmt="html"))
+        return f"<pre>{table}</pre>" + footer
+
+    def select(self, list_of_keys: List[str]) -> "Scenario":
         """Select a subset of keys from a scenario.
 
         :param list_of_keys: The keys to select.
@@ -270,7 +266,7 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
             new_scenario[key] = self[key]
         return new_scenario
 
-    def drop(self, list_of_keys: Collection[str]) -> "Scenario":
+    def drop(self, list_of_keys: List[str]) -> "Scenario":
         """Drop a subset of keys from a scenario.
 
         :param list_of_keys: The keys to drop.
@@ -324,7 +320,7 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         ...     _ = f.flush()
         ...     s = Scenario.from_file(f.name, "file")
         >>> s
-        Scenario({'file': FileStore(path='...', ...)})
+        Scenario({'file': FileStore(path='...')})
 
         """
         from edsl.scenarios.FileStore import FileStore
@@ -355,10 +351,35 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         return cls.from_file(image_path, image_name)
 
     @classmethod
-    def from_pdf(cls, pdf_path: str):
-        from edsl.scenarios.PdfExtractor import PdfExtractor
+    def from_pdf(cls, pdf_path):
+        # Ensure the file exists
+        import fitz
+
+        if not os.path.exists(pdf_path):
+            raise FileNotFoundError(f"The file {pdf_path} does not exist.")
+
+        # Open the PDF file
+        document = fitz.open(pdf_path)
+
+        # Get the filename from the path
+        filename = os.path.basename(pdf_path)
 
-        return PdfExtractor(pdf_path, cls).get_object()
+        # Iterate through each page and extract text
+        text = ""
+        for page_num in range(len(document)):
+            page = document.load_page(page_num)
+            blocks = page.get_text("blocks")  # Extract text blocks
+
+            # Sort blocks by their vertical position (y0) to maintain reading order
+            blocks.sort(key=lambda b: (b[1], b[0]))  # Sort by y0 first, then x0
+
+            # Combine the text blocks in order
+            for block in blocks:
+                text += block[4] + "\n"
+
+        # Create a dictionary for the combined text
+        page_info = {"filename": filename, "text": text}
+        return Scenario(page_info)
 
     @classmethod
     def from_docx(cls, docx_path: str) -> "Scenario":
@@ -378,9 +399,52 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         Scenario({'file_path': 'test.docx', 'text': 'EDSL Survey\\nThis is a test.'})
         >>> import os; os.remove("test.docx")
         """
-        from edsl.scenarios.DocxScenario import DocxScenario
+        from docx import Document
+
+        doc = Document(docx_path)
 
-        return Scenario(DocxScenario(docx_path).get_scenario_dict())
+        # Extract all text
+        full_text = []
+        for para in doc.paragraphs:
+            full_text.append(para.text)
+
+        # Join the text from all paragraphs
+        text = "\n".join(full_text)
+        return Scenario({"file_path": docx_path, "text": text})
+
+    @staticmethod
+    def _line_chunks(text, num_lines: int) -> Generator[str, None, None]:
+        """Split a text into chunks of a given size.
+
+        :param text: The text to split.
+        :param num_lines: The number of lines in each chunk.
+
+        Example:
+
+        >>> list(Scenario._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.']
+        """
+        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, num_words: int) -> Generator[str, None, None]:
+        """Split a text into chunks of a given size.
+
+        :param text: The text to split.
+        :param num_words: The number of words in each chunk.
+
+        Example:
+
+        >>> list(Scenario._word_chunks("This is a test.", 2))
+        ['This is', 'a test.']
+        """
+        words = text.split()
+        for i in range(0, len(words), num_words):
+            chunk = " ".join(words[i : i + num_words])
+            yield chunk
 
     def chunk(
         self,
@@ -431,11 +495,36 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         ...
         ValueError: You must specify either num_words or num_lines, but not both.
         """
-        from edsl.scenarios.DocumentChunker import DocumentChunker
+        from edsl.scenarios.ScenarioList import ScenarioList
 
-        return DocumentChunker(self).chunk(
-            field, num_words, num_lines, include_original, hash_original
-        )
+        if num_words is not None:
+            chunks = list(self._word_chunks(self[field], num_words))
+
+        if num_lines is not None:
+            chunks = list(self._line_chunks(self[field], num_lines))
+
+        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."
+            )
+
+        scenarios = []
+        for i, chunk in enumerate(chunks):
+            new_scenario = copy.deepcopy(self)
+            new_scenario[field] = chunk
+            new_scenario[field + "_chunk"] = i
+            if include_original:
+                if hash_original:
+                    new_scenario[field + "_original"] = hashlib.md5(
+                        self[field].encode()
+                    ).hexdigest()
+                else:
+                    new_scenario[field + "_original"] = self[field]
+            scenarios.append(new_scenario)
+        return ScenarioList(scenarios)
 
     @classmethod
     @remove_edsl_version
@@ -458,30 +547,44 @@ class Scenario(Base, UserDict, ScenarioHtmlMixin):
         return cls(d)
 
     def _table(self) -> tuple[dict, list]:
-        """Prepare generic table data.
-        >>> s = Scenario({"food": "wood chips"})
-        >>> s._table()
-        ([{'Attribute': 'data', 'Value': "{'food': 'wood chips'}"}, {'Attribute': 'name', 'Value': 'None'}], ['Attribute', 'Value'])
-        """
+        """Prepare generic table data."""
         table_data = []
         for attr_name, attr_value in self.__dict__.items():
             table_data.append({"Attribute": attr_name, "Value": repr(attr_value)})
         column_names = ["Attribute", "Value"]
         return table_data, column_names
 
+    def rich_print(self) -> "Table":
+        """Display an object as a rich table."""
+        from rich.table import Table
+
+        table_data, column_names = self._table()
+        table = Table(title=f"{self.__class__.__name__} Attributes")
+        for column in column_names:
+            table.add_column(column, style="bold")
+
+        for row in table_data:
+            row_data = [row[column] for column in column_names]
+            table.add_row(*row_data)
+
+        return table
+
     @classmethod
-    def example(cls, randomize: bool = False) -> Scenario:
+    def example(cls, randomize: bool = False, has_image=False) -> Scenario:
         """
         Returns an example Scenario instance.
 
         :param randomize: If True, adds a random string to the value of the example key.
         """
-        addition = "" if not randomize else str(uuid4())
-        return cls(
-            {
-                "persona": f"A reseacher studying whether LLMs can be used to generate surveys.{addition}",
-            }
-        )
+        if not has_image:
+            addition = "" if not randomize else str(uuid4())
+            return cls(
+                {
+                    "persona": f"A reseacher studying whether LLMs can be used to generate surveys.{addition}",
+                }
+            )
+        else:
+            return cls.from_image(cls.example_image())
 
     def code(self) -> List[str]:
         """Return the code for the scenario."""

edsl/scenarios/ScenarioHtmlMixin.py

@@ -1,4 +1,7 @@
+import requests
 from typing import Optional
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
 
 
 class ScenarioHtmlMixin:
@@ -19,10 +22,6 @@ class ScenarioHtmlMixin:
 
     def fetch_html(url):
         # Define the user-agent to mimic a browser
-        import requests
-        from requests.adapters import HTTPAdapter
-        from requests.packages.urllib3.util.retry import Retry
-
         headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
         }

edsl/scenarios/{scenario_join.py → ScenarioJoin.py}

@@ -1,9 +1,9 @@
 from __future__ import annotations
 from typing import Union, TYPE_CHECKING
 
-if TYPE_CHECKING:
-    from edsl.scenarios.ScenarioList import ScenarioList
-    from edsl.scenarios.Scenario import Scenario
+# if TYPE_CHECKING:
+from edsl.scenarios.ScenarioList import ScenarioList
+from edsl.scenarios.Scenario import Scenario
 
 
 class ScenarioJoin:
@@ -23,7 +23,7 @@ class ScenarioJoin:
         self.left = left
         self.right = right
 
-    def left_join(self, by: Union[str, list[str]]) -> "ScenarioList":
+    def left_join(self, by: Union[str, list[str]]) -> ScenarioList:
         """Perform a left join between the two ScenarioLists.
 
         Args:
@@ -35,8 +35,6 @@ class ScenarioJoin:
         Raises:
             ValueError: If by is empty or if any join keys don't exist in both ScenarioLists
         """
-        from edsl.scenarios.ScenarioList import ScenarioList
-
         self._validate_join_keys(by)
         by_keys = [by] if isinstance(by, str) else by
 
@@ -88,8 +86,6 @@ class ScenarioJoin:
         self, by_keys: list[str], other_dict: dict, all_keys: set
     ) -> list[Scenario]:
         """Create the joined scenarios."""
-        from edsl.scenarios.Scenario import Scenario
-
         new_scenarios = []
 
         for scenario in self.left:
@@ -109,8 +105,8 @@ class ScenarioJoin:
     def _handle_matching_scenario(
         self,
         new_scenario: dict,
-        left_scenario: "Scenario",
-        right_scenario: "Scenario",
+        left_scenario: Scenario,
+        right_scenario: Scenario,
         by_keys: list[str],
     ) -> None:
         """Handle merging of matching scenarios and conflict warnings."""