Make NoteableClient sync contextmanager (#42)

* Sync contextmanager for NoteableClient

* read cell state from bulk cell state update event

* fix CI

* add changelog

* add changelog

.github/workflows/ci.yaml

@@ -22,11 +22,11 @@ jobs:
       - name: Setup poetry
         uses: Gr1N/setup-poetry@v7
         with:
-          poetry-version: 1.2.0
+          poetry-version: 1.2.1
       - name: Install dependencies
         run: |
-          poetry export -f requirements.txt --extras "all" --without-hashes --with dev --output requirements.txt
-          pip install -r requirements.txt
+          poetry install --all-extras --with dev
+          pip install nox-poetry
       - name: Run tests
         run: nox -s test-${{ matrix.python_version }}
       - name: Generate coverage XML
@@ -43,10 +43,10 @@ jobs:
       - name: Setup poetry
         uses: Gr1N/setup-poetry@v7
         with:
-          poetry-version: 1.2.0
+          poetry-version: 1.2.1
       - name: Install dependencies
         run: |
-          poetry export -f requirements.txt --extras "all" --without-hashes --with dev --output requirements.txt
-          pip install -r requirements.txt
+          poetry install --all-extras --with dev
+          pip install nox-poetry
       - name: Run lint
         run: nox -s lint

CHANGELOG.md

@@ -5,6 +5,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
+### Added
+- Add `__enter__` and `__exit__` methods to `NoteableClient` to allow it to be used as a synchronous context manager
+- Read `bulk_cell_state_update_event` messages for cell state instead of `cell_state_update_event`
+- Add `rename` to `AccessLevelActions` enum
+
 ## [0.0.5] - 2022-09-13
 ### Added
 - Add kernel output types for papermill-origami

origami/client.py

@@ -16,8 +16,11 @@
 import structlog
 import websockets
 from httpx import ReadTimeout
+from nbclient.util import run_sync
 from pydantic import BaseModel, BaseSettings, ValidationError
 
+from origami.types.rtu import BulkCellStateMessage
+
 from .types.deltas import FileDeltaAction, FileDeltaType, NBCellProperties, V2CellContentsProperties
 from .types.files import FileVersion, NotebookFile
 from .types.jobs import (
@@ -34,7 +37,6 @@
     AuthenticationRequest,
     AuthenticationRequestData,
     CallbackTracker,
-    CellStateMessageReply,
     FileDeltaReply,
     FileSubscribeReplySchema,
     GenericRTUMessage,
@@ -130,7 +132,7 @@ def __init__(
         if not config:
             settings = ClientSettings()
             if not os.path.exists(settings.auth0_config_path):
-                logger.error(
+                logger.warning(
                     f"No config object passed in and no config file found at {settings.auth0_config_path}"
                     ", using default empty config"
                 )
@@ -409,6 +411,12 @@ async def __aexit__(self, exc_type, exc, tb):
         finally:
             return await httpx.AsyncClient.__aexit__(self, exc_type, exc, tb)
 
+    def __enter__(self):
+        return run_sync(self.__aenter__)()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return run_sync(self.__aexit__)(exc_type, exc_val, exc_tb)
+
     def register_message_callback(
         self,
         callable: RTUEventCallable,
@@ -784,12 +792,13 @@ async def check_success(resp: GenericRTUReply):
         tracker_future = tracker.next_trigger
         results_tracker_future = None
 
-        async def cell_complete_check(resp: CellStateMessageReply):
-            if resp.data.cell_id != cell_id:
-                raise SkipCallback("Not tracked cell")
-            if not resp.data.state.is_terminal_state:
-                raise SkipCallback("Not terminal state")
-            return resp
+        async def cell_complete_check(resp: BulkCellStateMessage):
+            for cell_state in resp.data.cell_states:
+                if cell_state.cell_id == cell_id and cell_state.state.is_terminal_state:
+                    return cell_state
+                elif cell_state.cell_id == cell_id and not cell_state.state.is_terminal_state:
+                    raise SkipCallback("Not terminal state")
+            raise SkipCallback("Not tracked cell")
 
         if await_results:
             assert (
@@ -799,8 +808,8 @@ async def cell_complete_check(resp: CellStateMessageReply):
             results_tracker = self.register_message_callback(
                 cell_complete_check,
                 session.kernel_channel,
-                "cell_state_update_event",
-                response_schema=CellStateMessageReply,
+                "bulk_cell_state_update_event",
+                response_schema=BulkCellStateMessage,
             )
             results_tracker_future = results_tracker.next_trigger
 

origami/types/access_levels.py

@@ -65,6 +65,7 @@ def _generate_next_value_(name, start, count, last_values):
     create_file_version = enum.auto()
     create_file_sandbox = enum.auto()
     create_parameterized_notebook = enum.auto()
+    rename = enum.auto()
 
     # NotebookFile comment actions
     view_comments = enum.auto()

origami/types/rtu.py

@@ -314,10 +314,15 @@ class CellStateMessageData(BaseModel):
     queued_by_id: Optional[UUID]  # user ID
 
 
-class CellStateMessageReply(GenericRTUReplySchema[CellStateMessageData]):
-    """Defines a status update message from the rtu websocket"""
+class BulkCellStateMessageData(BaseModel):
 
-    event = 'cell_state_update_event'
+    cell_states: List[CellStateMessageData]
+
+
+class BulkCellStateMessage(GenericRTUReplySchema[BulkCellStateMessageData]):
+    """Defines a bulk cell status update message from the rtu websocket"""
+
+    event = 'bulk_cell_state_update_event'
 
 
 class KernelOutputType(enum.Enum):