edsl 0.1.27.dev2__py3-none-any.whl → 0.1.29__py3-none-any.whl

edsl/notebooks/Notebook.py

@@ -1,7 +1,9 @@
-"""A Notebook is ...."""
+"""A Notebook is a utility class that allows you to easily share/pull ipynbs from Coop."""
 
+import json
 from typing import Dict, List, Optional
-from rich.table import Table
+
+
 from edsl.Base import Base
 from edsl.utilities.decorators import (
     add_edsl_version,
@@ -14,43 +16,91 @@ class Notebook(Base):
     A Notebook is a utility class that allows you to easily share/pull ipynbs from Coop.
     """
 
-    def __init__(self, data: Optional[Dict] = None, path: Optional[str] = None):
+    default_name = "notebook"
+
+    def __init__(
+        self,
+        data: Optional[Dict] = None,
+        path: Optional[str] = None,
+        name: Optional[str] = None,
+    ):
         """
         Initialize a new Notebook.
-        - if a path is provided, try to load the notebook from that path.
-        - if no path is provided, assume this code is run in a notebook and try to load the current notebook.
+
+        :param data: A dictionary representing the notebook data.
+        This dictionary must conform to the official Jupyter Notebook format, as defined by nbformat.
+        :param path: A filepath from which to load the notebook.
+        If no path is provided, assume this code is run in a notebook and try to load the current notebook from file.
+        :param name: A name for the Notebook.
         """
+        import nbformat
+
+        # Load current notebook path as fallback (VS Code only)
+        path = path or globals().get("__vsc_ipynb_file__")
         if data is not None:
+            nbformat.validate(data)
             self.data = data
         elif path is not None:
-            # TO BE IMPLEMENTED
-            # store in this var the data from the notebook
-            self.data = {"some": "data"}
+            with open(path, mode="r", encoding="utf-8") as f:
+                data = nbformat.read(f, as_version=4)
+            self.data = json.loads(json.dumps(data))
         else:
-            # TO BE IMPLEMENTED
-            # 1. Check you're in a notebook ...
-            # 2. get its info and store it in self.data
-            self.data = {"some": "data"}
+            # TODO: Support for IDEs other than VSCode
+            raise NotImplementedError(
+                "Cannot create a notebook from within itself in this development environment"
+            )
 
-        # deprioritize - perhaps add sanity check function
+        # TODO: perhaps add sanity check function
         # 1. could check if the notebook is a valid notebook
         # 2. could check notebook uses EDSL
         # ....
 
+        self.name = name or self.default_name
+
+    @classmethod
+    def from_script(cls, path: str, name: Optional[str] = None) -> "Notebook":
+        # Read the script file
+        with open(path, "r") as script_file:
+            script_content = script_file.read()
+
+        # Create a new Jupyter notebook
+        nb = nbformat.v4.new_notebook()
+
+        # Add the script content to the first cell
+        first_cell = nbformat.v4.new_code_cell(script_content)
+        nb.cells.append(first_cell)
+
+        # Create a Notebook instance with the notebook data
+        notebook_instance = cls(nb)
+
+        return notebook_instance
+
+    @classmethod
+    def from_current_script(cls) -> "Notebook":
+        import inspect
+        import os
+
+        # Get the path to the current file
+        current_frame = inspect.currentframe()
+        caller_frame = inspect.getouterframes(current_frame, 2)
+        current_file_path = os.path.abspath(caller_frame[1].filename)
+
+        # Use from_script to create the notebook
+        return cls.from_script(current_file_path)
+
     def __eq__(self, other):
         """
         Check if two Notebooks are equal.
+        This only checks the notebook data.
         """
         return self.data == other.data
 
     @add_edsl_version
     def to_dict(self) -> dict:
         """
-        Convert to a dictionary.
-        AF: here you will create a dict from which self.from_dict can recreate the object.
-        AF: the decorator will add the edsl_version to the dict.
+        Convert a Notebook to a dictionary.
         """
-        return {"data": self.data}
+        return {"name": self.name, "data": self.data}
 
     @classmethod
     @remove_edsl_version
@@ -58,12 +108,17 @@ class Notebook(Base):
         """
         Convert a dictionary representation of a Notebook to a Notebook object.
         """
-        return cls(data=d["data"])
+        return cls(data=d["data"], name=d["name"])
+
+    def to_file(self, path: str):
+        """
+        Save the notebook at the specified filepath.
+        """
+        nbformat.write(nbformat.from_dict(self.data), fp=path)
 
     def print(self):
         """
         Print the notebook.
-        AF: not sure how this should behave for a notebook
         """
         from rich import print_json
         import json
@@ -72,45 +127,115 @@ class Notebook(Base):
 
     def __repr__(self):
         """
-        AF: not sure how this should behave for a notebook
+        Return representation of Notebook.
         """
-        return f"Notebook({self.to_dict()})"
+        return f'Notebook(data={self.data}, name="""{self.name}""")'
 
     def _repr_html_(self):
         """
-        AF: not sure how this should behave for a notebook
+        Return HTML representation of Notebook.
+        """
+        from nbconvert import HTMLExporter
+        import nbformat
+
+        notebook = nbformat.from_dict(self.data)
+        html_exporter = HTMLExporter(template_name="basic")
+        (body, _) = html_exporter.from_notebook_node(notebook)
+        return body
+
+    def _table(self) -> tuple[dict, list]:
+        """
+        Prepare generic table data.
         """
-        from edsl.utilities.utilities import data_to_html
+        table_data = []
+
+        notebook_preview = ""
+        for cell in self.data["cells"]:
+            if "source" in cell:
+                notebook_preview += f"{cell['source']}\n"
+            if len(notebook_preview) > 1000:
+                notebook_preview = f"{notebook_preview[:1000]} [...]"
+                break
+        notebook_preview = notebook_preview.rstrip()
+
+        table_data.append(
+            {
+                "Attribute": "name",
+                "Value": repr(self.name),
+            }
+        )
+        table_data.append(
+            {
+                "Attribute": "notebook_preview",
+                "Value": notebook_preview,
+            }
+        )
 
-        return data_to_html(self.to_dict())
+        column_names = ["Attribute", "Value"]
+        return table_data, column_names
 
     def rich_print(self) -> "Table":
         """
-        AF: not sure how we should implement this for a notebook
+        Display a Notebook as a rich table.
         """
-        pass
+        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) -> "Notebook":
         """
         Return an example Notebook.
-        AF: add a simple custom example here
         """
-        return cls(data={"some": "data"})
+        cells = [
+            {
+                "cell_type": "markdown",
+                "metadata": dict(),
+                "source": "# Test notebook",
+            },
+            {
+                "cell_type": "code",
+                "execution_count": 1,
+                "metadata": dict(),
+                "outputs": [
+                    {
+                        "name": "stdout",
+                        "output_type": "stream",
+                        "text": "Hello world!\n",
+                    }
+                ],
+                "source": 'print("Hello world!")',
+            },
+        ]
+        data = {
+            "metadata": dict(),
+            "nbformat": 4,
+            "nbformat_minor": 4,
+            "cells": cells,
+        }
+        return cls(data=data)
 
     def code(self) -> List[str]:
         """
         Return the code that could be used to create this Notebook.
-        AF: Again, not sure
         """
         lines = []
-        lines.append("from edsl.notebooks import Notebook")
-        lines.append(f"s = Notebook({self.data})")
+        lines.append("from edsl import Notebook")
+        lines.append(f'nb = Notebook(data={self.data}, name="""{self.name}""")')
         return lines
 
 
 if __name__ == "__main__":
-    from edsl.notebooks import Notebook
+    from edsl import Notebook
 
     notebook = Notebook.example()
     assert notebook == notebook.from_dict(notebook.to_dict())

edsl/prompts/Prompt.py

@@ -1,12 +1,16 @@
-"""Class for creating prompts to be used in a survey."""
-
 from __future__ import annotations
 from typing import Optional
 from abc import ABC
 from typing import Any, List
 
 from rich.table import Table
-from jinja2 import Template, Environment, meta, TemplateSyntaxError
+from jinja2 import Template, Environment, meta, TemplateSyntaxError, Undefined
+
+
+class PreserveUndefined(Undefined):
+    def __str__(self):
+        return "{{ " + self._undefined_name + " }}"
+
 
 from edsl.exceptions.prompts import TemplateRenderError
 from edsl.prompts.prompt_config import (
@@ -35,6 +39,10 @@ class PromptBase(
 
         return data_to_html(self.to_dict())
 
+    def __len__(self):
+        """Return the length of the prompt text."""
+        return len(self.text)
+
     @classmethod
     def prompt_attributes(cls) -> List[str]:
         """Return the prompt class attributes."""
@@ -75,10 +83,10 @@ class PromptBase(
         >>> p = Prompt("Hello, {{person}}")
         >>> p2 = Prompt("How are you?")
         >>> p + p2
-        Prompt(text='Hello, {{person}}How are you?')
+        Prompt(text=\"""Hello, {{person}}How are you?\""")
 
         >>> p + "How are you?"
-        Prompt(text='Hello, {{person}}How are you?')
+        Prompt(text=\"""Hello, {{person}}How are you?\""")
         """
         if isinstance(other_prompt, str):
             return self.__class__(self.text + other_prompt)
@@ -114,7 +122,7 @@ class PromptBase(
         Example:
         >>> p = Prompt("Hello, {{person}}")
         >>> p
-        Prompt(text='Hello, {{person}}')
+        Prompt(text=\"""Hello, {{person}}\""")
         """
         return f'Prompt(text="""{self.text}""")'
 
@@ -137,7 +145,7 @@ class PromptBase(
         :param template: The template to find the variables in.
 
         """
-        env = Environment()
+        env = Environment(undefined=PreserveUndefined)
         ast = env.parse(template)
         return list(meta.find_undeclared_variables(ast))
 
@@ -186,13 +194,16 @@ class PromptBase(
 
         >>> p = Prompt("Hello, {{person}}")
         >>> p.render({"person": "John"})
-        'Hello, John'
+        Prompt(text=\"""Hello, John\""")
 
         >>> p.render({"person": "Mr. {{last_name}}", "last_name": "Horton"})
-        'Hello, Mr. Horton'
+        Prompt(text=\"""Hello, Mr. Horton\""")
 
         >>> p.render({"person": "Mr. {{last_name}}", "last_name": "Ho{{letter}}ton"}, max_nesting = 1)
-        'Hello, Mr. Horton'
+        Prompt(text=\"""Hello, Mr. Ho{{ letter }}ton\""")
+
+        >>> p.render({"person": "Mr. {{last_name}}"})
+        Prompt(text=\"""Hello, Mr. {{ last_name }}\""")
         """
         new_text = self._render(
             self.text, primary_replacement, **additional_replacements
@@ -216,12 +227,13 @@ class PromptBase(
         >>> codebook = {"age": "Age"}
         >>> p = Prompt("You are an agent named {{ name }}. {{ codebook['age']}}: {{ age }}")
         >>> p.render({"name": "John", "age": 44}, codebook=codebook)
-        'You are an agent named John. Age: 44'
+        Prompt(text=\"""You are an agent named John. Age: 44\""")
         """
+        env = Environment(undefined=PreserveUndefined)
         try:
             previous_text = None
             for _ in range(MAX_NESTING):
-                rendered_text = Template(text).render(
+                rendered_text = env.from_string(text).render(
                     primary_replacement, **additional_replacements
                 )
                 if rendered_text == previous_text:
@@ -258,7 +270,7 @@ class PromptBase(
         >>> p = Prompt("Hello, {{person}}")
         >>> p2 = Prompt.from_dict(p.to_dict())
         >>> p2
-        Prompt(text='Hello, {{person}}')
+        Prompt(text=\"""Hello, {{person}}\""")
 
         """
         class_name = data["class_name"]
@@ -290,6 +302,12 @@ class Prompt(PromptBase):
     component_type = ComponentTypes.GENERIC
 
 
+if __name__ == "__main__":
+    print("Running doctests...")
+    import doctest
+
+    doctest.testmod()
+
 from edsl.prompts.library.question_multiple_choice import *
 from edsl.prompts.library.agent_instructions import *
 from edsl.prompts.library.agent_persona import *
@@ -302,9 +320,3 @@ from edsl.prompts.library.question_numerical import *
 from edsl.prompts.library.question_rank import *
 from edsl.prompts.library.question_extract import *
 from edsl.prompts.library.question_list import *
-
-
-if __name__ == "__main__":
-    import doctest
-
-    doctest.testmod()