Merge pull request #210 from pmelchior/sean/questionnaire_start

Questionnaire UI and initial functionality

Author: smcguire-cmu

pyproject.toml

@@ -26,9 +26,13 @@ dependencies = [
     "optax",
     "numpyro",
     "h5py",
+    "ipywidgets",
     "pydantic",
     "pyyaml",
     "colorama",
+    "markdown",
+    "pygments",
+    "jinja2"
 ]
 
 [project.urls]
@@ -43,6 +47,7 @@ dev = [
     "pre-commit", # Used to run checks before finalizing a git commit
     "pytest",
     "pytest-cov", # Used to report total code coverage
+    "pytest-mock", # Used to mock objects in tests
     "ruff", # Used for static linting of files
 ]
 
@@ -56,9 +61,6 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools_scm]
 write_to = "src/scarlet2/_version.py"
 
-[tool.setuptools.package-data]
-"scarlet2.questionnaire" = ["questions.yaml"]
-
 [tool.pytest.ini_options]
 testpaths = [
     "tests",

src/scarlet2/questionnaire/__init__.py

@@ -0,0 +1,3 @@
+from .questionnaire import QuestionnaireWidget, run_questionnaire
+
+__all__ = ["QuestionnaireWidget", "run_questionnaire"]

src/scarlet2/questionnaire/models.py

@@ -34,4 +34,5 @@ class Questionnaire(BaseModel):
     """Represents a questionnaire with an initial template and a list of questions."""
 
     initial_template: str
+    initial_commentary: str = ""
     questions: list[Question]

src/scarlet2/questionnaire/questionnaire.py

@@ -1,11 +1,156 @@
+import re
 from importlib.resources import files
 
+import jinja2
+import markdown
 import yaml
+from IPython.display import display
+from ipywidgets import HTML, Button, HBox, Label, Layout, VBox
+from pygments import highlight
+from pygments.formatters import HtmlFormatter
+from pygments.lexers import PythonLexer
 
-from scarlet2.questionnaire.models import Questionnaire
+from scarlet2.questionnaire.models import Question, Questionnaire, Template
 
-FILE_PACKAGE_PATH = "scarlet2.questionnaire"
-FILE_NAME = "questions.yaml"
+PACKAGE_PATH = "scarlet2.questionnaire"
+QUESTIONS_FILE_NAME = "questions.yaml"
+
+VIEWS_PACKAGE_PATH = "scarlet2.questionnaire.views"
+OUTPUT_BOX_TEMPLATE_FILE = "output_box.html.jinja"
+OUTPUT_BOX_STYLE_FILE = "output_box.css"
+
+
+QUESTION_BOX_LAYOUT = Layout(
+    padding="12px",
+    backgroundColor="#f9f9f9",
+    border="1px solid #ddd",
+    borderRadius="10px",
+    width="45%",
+)
+
+OUTPUT_BOX_LAYOUT = Layout(
+    width="50%",
+    margin="0 0 0 20px",
+)
+
+
+class QuestionnaireWidget:
+    """A widget to run an interactive questionnaire in a Jupyter notebook."""
+
+    def __init__(self, questionnaire: Questionnaire):
+        self.questions = questionnaire.questions
+        self.code_output = questionnaire.initial_template
+        self.commentary = questionnaire.initial_commentary
+
+        self._init_questions()
+        self._init_ui()
+
+        self._render_output_box()
+        self._render_next_question()
+
+    def _init_questions(self):
+        self.current_question = None
+        self.questions_stack = []
+        self.question_answers = []
+        self._add_questions_to_stack(self.questions)
+
+    def _init_ui(self):
+        self.output_container = HTML()
+        self.question_box = VBox([], layout=QUESTION_BOX_LAYOUT)
+        self.output_box = VBox([self.output_container], layout=OUTPUT_BOX_LAYOUT)
+
+        self.ui = HBox([self.question_box, self.output_box])
+
+    def _add_questions_to_stack(self, questions: list[Question]):
+        self.questions_stack = questions + self.questions_stack
+
+    def _render_next_question(self):
+        self.current_question = self.questions_stack.pop(0) if len(self.questions_stack) > 0 else None
+        self._render_question_box()
+
+    def _render_question_box(self):
+        previous_qs = self._generate_previous_questions()
+
+        if self.current_question is None:
+            self.question_box.children = previous_qs + [Label("🎉 You're done!")]
+            return
+
+        q_label = HTML(f"<b>{self.current_question.question}</b>")
+
+        buttons = []
+        for i, answer in enumerate(self.current_question.answers):
+            button = Button(
+                description=answer.answer,
+                tooltip=answer.tooltip,
+                layout=Layout(width="auto", margin="4px 0"),
+            )
+
+            def on_click_handler(btn, index=i):
+                self._handle_answer(index)
+
+            button.on_click(on_click_handler)
+
+            buttons.append(button)
+
+        self.question_box.children = previous_qs + [q_label] + buttons
+
+    def _generate_previous_questions(self):
+        children = []
+        for q, ans_ind in self.question_answers:
+            html_str = f"""
+            <div style="background_color: #111">
+                <span style="color: #888; padding-right: 10px;">{q.question}</span>
+                <span style="color: #555;">{q.answers[ans_ind].answer}</span>
+            </div>
+            """
+            children.append(HTML(html_str))
+        return children
+
+    def _handle_answer(self, answer_index: int):
+        answer = self.current_question.answers[answer_index]
+
+        self._update_template(answer.templates)
+        self.commentary = answer.commentary
+        self._render_output_box()
+
+        self._add_questions_to_stack(answer.followups)
+        self.question_answers.append((self.current_question, answer_index))
+        self._render_next_question()
+
+    def _update_template(self, templates: list[Template]):
+        for t in templates:
+            pattern = r"\{\{\s*" + re.escape(t.replacement) + r"\s*\}\}"
+            self.code_output = re.sub(pattern, t.code, self.code_output)
+
+    def _render_output_box(self):
+        output_code = re.sub(r"\{\{\s*\w+\s*\}\}", "", self.code_output)
+        commentary_html = markdown.markdown(self.commentary, extensions=["extra"])
+
+        formatter = HtmlFormatter(style="monokai", noclasses=True)
+        highlighted_code = highlight(output_code, PythonLexer(), formatter)
+
+        # Load Jinja template
+        template_file = files(VIEWS_PACKAGE_PATH).joinpath(OUTPUT_BOX_TEMPLATE_FILE)
+        with template_file.open("r") as f:
+            template_source = f.read()
+        template = jinja2.Template(template_source)
+
+        # Load additional CSS for styling
+        css_file = files(VIEWS_PACKAGE_PATH).joinpath(OUTPUT_BOX_STYLE_FILE)
+        with css_file.open("r") as f:
+            css_content = f.read()
+
+        html_content = f"<style>{css_content}</style>\n" + template.render(
+            highlighted_code=highlighted_code,
+            raw_code=output_code,
+            commentary_html=commentary_html,
+        )
+
+        self.output_container.value = html_content
+
+    def show(self):
+        """Display the widget in a Jupyter notebook."""
+        display(self.ui)
 
 
 def load_questions() -> Questionnaire:
@@ -14,7 +159,22 @@ def load_questions() -> Questionnaire:
     Returns:
         Questionnaire: The loaded questionnaire model.
     """
-    questions_path = files(FILE_PACKAGE_PATH).joinpath(FILE_NAME)
+    questions_path = files(PACKAGE_PATH).joinpath(QUESTIONS_FILE_NAME)
     with questions_path.open("r") as f:
         raw = yaml.safe_load(f)
         return Questionnaire.model_validate(raw)
+
+
+def run_questionnaire():
+    """Run the Scarlet2 initialization questionnaire in a Jupyter notebook.
+
+    The questionnaire guides the user through a series of questions to set up
+    the initialization of a Scarlet2 project that fits their use case.
+
+    The user will be presented with questions and multiple-choice answers, and
+    at the end of the questionnaire, a code snippet that can be used as a
+    template for initializing Scarlet2 will be generated.
+    """
+    questions = load_questions()
+    app = QuestionnaireWidget(questions)
+    app.show()

src/scarlet2/questionnaire/views/__init__.py

@@ -0,0 +1,29 @@
+.code-output {
+    background-color: #272822;
+    color: #f1f1f1;
+    padding: 10px;
+    border-radius: 10px;
+    border: 1px solid #444;
+    font-family: monospace;
+    overflow-x: auto;
+}
+
+.copy_button {
+    font-size: 12px;
+    padding: 4px 8px;
+    background-color: #272822;
+    color: white;
+    border: none;
+    border-radius: 4px;
+    cursor: pointer;
+}
+
+.commentary-box {
+    background-color: #f6f8fa;
+    color: #333;
+    padding: 8px 12px;
+    border-left: 4px solid #0366d6;
+    border-radius: 6px;
+    font-size: 14px;
+    font-family: sans-serif;
+}

src/scarlet2/questionnaire/views/output_box.css

@@ -0,0 +1,12 @@
+<div class="code-output">
+    {{highlighted_code | safe}}
+    <div style="display: flex; justify-content: flex-end; margin-top: 8px;">
+        <button onclick='navigator.clipboard.writeText({{ raw_code | tojson | safe }})'
+                class="copy_button">
+            📋 Copy
+        </button>
+    </div>
+</div>
+<div class="commentary-box">
+    {{commentary_html | safe}}
+</div>