Convert all models to use Annotated

lsst-sqre · Mar 6, 2024 · 0a44b51 · 0a44b51
1 parent e113c69
commit 0a44b51
Show file tree

Hide file tree

Showing 5 changed files with 201 additions and 154 deletions.
diff --git a/src/noteburst/handlers/external.py b/src/noteburst/handlers/external.py
@@ -20,7 +20,7 @@
 class Index(BaseModel):
     """Metadata about the application."""
 
-    metadata: SafirMetadata = Field(..., title="Package metadata")
+    metadata: Annotated[SafirMetadata, Field(title="Package metadata")]
 
 
 @external_router.get(

diff --git a/src/noteburst/handlers/v1/models.py b/src/noteburst/handlers/v1/models.py
@@ -4,7 +4,7 @@
 
 import json
 from datetime import datetime
-from typing import Any
+from typing import Annotated, Any
 
 from arq.jobs import JobStatus
 from fastapi import Request
@@ -31,8 +31,8 @@
 class NotebookError(BaseModel):
     """Information about an exception that occurred during notebook exec."""
 
-    name: str = Field(description="The name of the exception.")
-    message: str = Field(description="The exception's message.")
+    name: Annotated[str, Field(description="The name of the exception.")]
+    message: Annotated[str, Field(description="The exception's message.")]
 
     @classmethod
     def from_nbexec_error(
@@ -52,56 +52,72 @@ class NotebookResponse(BaseModel):
     result and source notebooks.
     """
 
-    job_id: str = Field(title="The job ID")
-
-    kernel_name: str = kernel_name_field
-
-    enqueue_time: datetime = Field(
-        title="Time when the job was added to the queue (UTC)"
-    )
-
-    status: JobStatus = Field(
-        title="The current status of the notebook execution job"
-    )
-
-    self_url: AnyHttpUrl = Field(title="The URL of this resource")
-
-    source: str | None = Field(
-        None,
-        title="The content of the source ipynb file (JSON-encoded string)",
-        description="This field is null unless the source is requested.",
-    )
-
-    start_time: datetime | None = Field(
-        None,
-        title="Time when the notebook execution started (UTC)",
-        description="This field is present if the result is available.",
-    )
-
-    finish_time: datetime | None = Field(
-        None,
-        title="Time when the notebook execution completed (UTC)",
-        description="This field is present only if the result is available.",
-    )
-
-    success: bool | None = Field(
-        None,
-        title="Whether the execution was successful or not",
-        description="This field is present if the result is available.",
-    )
-
-    ipynb: str | None = Field(
-        None,
-        title="The contents of the executed Jupyter notebook",
-        description="The ipynb is a JSON-encoded string. This field is "
-        "present if the result is available.",
-    )
-
-    ipynb_error: NotebookError | None = Field(
-        None,
-        title="The error that occurred during notebook execution",
-        description="This field is null if an exeception did not occur.",
-    )
+    job_id: Annotated[str, Field(title="The job ID")]
+
+    kernel_name: Annotated[str, kernel_name_field]
+
+    enqueue_time: Annotated[
+        datetime, Field(title="Time when the job was added to the queue (UTC)")
+    ]
+
+    status: Annotated[
+        JobStatus,
+        Field(title="The current status of the notebook execution job"),
+    ]
+
+    self_url: Annotated[AnyHttpUrl, Field(title="The URL of this resource")]
+
+    source: Annotated[
+        str | None,
+        Field(
+            title="The content of the source ipynb file (JSON-encoded string)",
+            description="This field is null unless the source is requested.",
+        ),
+    ] = None
+
+    start_time: Annotated[
+        datetime | None,
+        Field(
+            title="Time when the notebook execution started (UTC)",
+            description="This field is present if the result is available.",
+        ),
+    ] = None
+
+    finish_time: Annotated[
+        datetime | None,
+        Field(
+            title="Time when the notebook execution completed (UTC)",
+            description=(
+                "This field is present only if the result is available."
+            ),
+        ),
+    ] = None
+
+    success: Annotated[
+        bool | None,
+        Field(
+            title="Whether the execution was successful or not",
+            description="This field is present if the result is available.",
+        ),
+    ] = None
+
+    ipynb: Annotated[
+        str | None,
+        Field(
+            title="The contents of the executed Jupyter notebook",
+            description="The ipynb is a JSON-encoded string. This field is "
+            "present if the result is available.",
+        ),
+    ] = None
+
+    ipynb_error: Annotated[
+        NotebookError | None,
+        Field(
+            None,
+            title="The error that occurred during notebook execution",
+            description="This field is null if an exeception did not occur.",
+        ),
+    ] = None
 
     @classmethod
     async def from_job_metadata(
@@ -112,6 +128,7 @@ async def from_job_metadata(
         include_source: bool = False,
         job_result: JobResult | None = None,
     ) -> NotebookResponse:
+        """Create a NotebookResponse from a job."""
         if job_result is not None and job_result.success:
             nbexec_result = NotebookExecutionResult.model_validate_json(
                 job_result.result
@@ -143,30 +160,35 @@ async def from_job_metadata(
 class PostNotebookRequest(BaseModel):
     """The ``POST /notebooks/`` request body."""
 
-    ipynb: str | dict[str, Any] = Field(
-        ...,
-        title="The contents of a Jupyter notebook",
-        description="If a string, the content is parsed as JSON. "
-        "Alternatively, the content can be submitted pre-parsed as "
-        "an object.",
-    )
-
-    kernel_name: str = kernel_name_field
-
-    enable_retry: bool = Field(
-        True,
-        title="Enable retries on failures",
-        description=(
-            "If true (default), noteburst will retry notebook "
-            "execution if the notebook fails, with an increasing back-off "
-            "time between tries. This is useful for dealing with transient "
-            "issues. However, if you are using Noteburst for continuous "
-            "integration of notebooks, disabling retries provides faster "
-            "feedback."
+    ipynb: Annotated[
+        str | dict[str, Any],
+        Field(
+            title="The contents of a Jupyter notebook",
+            description="If a string, the content is parsed as JSON. "
+            "Alternatively, the content can be submitted pre-parsed as "
+            "an object.",
+        ),
+    ]
+
+    kernel_name: Annotated[str, kernel_name_field]
+
+    enable_retry: Annotated[
+        bool,
+        Field(
+            title="Enable retries on failures",
+            description=(
+                "If true (default), noteburst will retry notebook "
+                "execution if the notebook fails, with an increasing back-off "
+                "time between tries. This is useful for dealing with "
+                "transient issues. However, if you are using Noteburst for "
+                "continuous integration of notebooks, disabling retries "
+                "provides faster feedback."
+            ),
         ),
-    )
+    ] = True
 
     def get_ipynb_as_str(self) -> str:
+        """Get the ipynb as a JSON-encoded string."""
         if isinstance(self.ipynb, str):
             return self.ipynb
         else:

diff --git a/src/noteburst/jupyterclient/jupyterlab.py b/src/noteburst/jupyterclient/jupyterlab.py
@@ -10,7 +10,7 @@
 from collections.abc import AsyncGenerator, AsyncIterator
 from dataclasses import dataclass
 from random import SystemRandom
-from typing import Any, Self
+from typing import Annotated, Any, Self
 from urllib.parse import urljoin, urlparse
 from uuid import uuid4
 
@@ -356,33 +356,38 @@ def __str__(self) -> str:
 
 
 class NotebookExecutionErrorModel(BaseModel):
-    """The error from the /user/:username/rubin/execute endpoint."""
+    """The error from the ``/user/:username/rubin/execute`` endpoint."""
 
-    traceback: str = Field(description="The exeception traceback.")
+    traceback: Annotated[str, Field(description="The exeception traceback.")]
 
-    ename: str = Field(description="The exception name.")
+    ename: Annotated[str, Field(description="The exception name.")]
 
-    evalue: str = Field(description="The exception value.")
+    evalue: Annotated[str, Field(description="The exception value.")]
 
-    err_msg: str = Field(description="The exception message.")
+    err_msg: Annotated[str, Field(description="The exception message.")]
 
 
 class NotebookExecutionResult(BaseModel):
     """The result of the /user/:username/rubin/execute endpoint."""
 
-    notebook: str = Field(
-        description="The notebook that was executed, as a JSON string."
-    )
+    notebook: Annotated[
+        str,
+        Field(description="The notebook that was executed, as a JSON string."),
+    ]
 
-    resources: dict[str, Any] = Field(
-        description=(
-            "The resources used to execute the notebook, as a JSON string."
-        )
-    )
+    resources: Annotated[
+        dict[str, Any],
+        Field(
+            description=(
+                "The resources used to execute the notebook, as a JSON string."
+            )
+        ),
+    ]
 
-    error: NotebookExecutionErrorModel | None = Field(
-        None, description="The error that occurred during execution."
-    )
+    error: Annotated[
+        NotebookExecutionErrorModel | None,
+        Field(description="The error that occurred during execution."),
+    ] = None
 
 
 class JupyterClient: