Refactor session extensions to use EventAwareExtension base class#8678
Merged
Refactor session extensions to use EventAwareExtension base class#8678
Conversation
- Updated various session extensions (e.g., ScratchCellListener, RecentsTrackerListener, CachingExtension) to inherit from EventAwareExtension, enabling automatic subscription to the event bus. - Introduced ExtensionRegistry for managing session extensions, replacing direct list manipulations. - Refactored event emission methods in SessionEventBus for improved clarity and error handling. - Adjusted related tests to ensure compatibility with the new extension structure. This change enhances the modularity and maintainability of session extensions.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors Marimo session extensions to reduce boilerplate and improve safety by centralizing event-bus subscription logic and introducing a small extensions registry abstraction.
Changes:
- Introduce
EventAwareExtension(auto subscribe/unsubscribe) andExtensionRegistry(typed lookup + snapshot iteration). - Refactor
SessionEventBusemission to use shared_emit/_emit_asynchelpers with listener-list snapshotting. - Update multiple extensions and tests to align with the new base class/registry behavior.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/_session/extensions/test_extensions.py | Updates lifecycle assertions and adds coverage for EventAwareExtension, ExtensionRegistry, and emit snapshot safety. |
| marimo/_session/session.py | Switches extensions storage to ExtensionRegistry and simplifies notification flush via typed lookup. |
| marimo/_session/file_watcher_integration.py | Migrates file watcher extension to EventAwareExtension and simplifies attach/detach logic. |
| marimo/_session/extensions/types.py | Adds EventAwareExtension + ExtensionRegistry definitions. |
| marimo/_session/extensions/extensions.py | Migrates core session extensions to EventAwareExtension; adds NotificationListenerExtension.flush(). |
| marimo/_session/events.py | Refactors event dispatch with snapshot iteration and shared error handling. |
| marimo/_server/session/listeners.py | Migrates RecentsTrackerListener to EventAwareExtension. |
| marimo/_server/scratchpad.py | Migrates ScratchCellListener to EventAwareExtension. |
| marimo/_internal/session/extensions.py | Re-exports new extension types in the internal API. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Add EventAwareExtension and ExtensionRegistry to the snapshot.
Use _canonicalize_path in on_attach to match on_detach and on_session_notebook_renamed, preventing potential watcher leaks when the raw path string differs from the normalized one.
manzt
previously approved these changes
Mar 30, 2026
Collaborator
manzt
left a comment
manzt
left a comment
There was a problem hiding this comment.
Some good DRY-ness! Nice.
marimo/_session/extensions/types.py
Outdated
| from marimo._session.events import SessionEventBus | ||
| from marimo._session.session import Session | ||
|
|
||
| _T = TypeVar("_T") |
Collaborator
There was a problem hiding this comment.
I think this is bounded by being a SessionExtension:
Suggested change
| _T = TypeVar("_T") | |
| _T = TypeVar("_T", bound=SessionExtension) |
marimo/_session/extensions/types.py
Outdated
|
|
||
| @property | ||
| def session(self) -> Session: | ||
| assert self._session is not None |
Collaborator
There was a problem hiding this comment.
Suggested change
| assert self._session is not None | |
| if self._session is None: | |
| raise RuntimeError("Extension is not attached to a session") |
marimo/_session/extensions/types.py
Outdated
|
|
||
| @property | ||
| def event_bus(self) -> SessionEventBus: | ||
| assert self._event_bus is not None |
Collaborator
There was a problem hiding this comment.
Suggested change
| assert self._event_bus is not None | |
| if self._event_bus is None: | |
| raise RuntimeError("Extension is not attached to a session") |
Comment on lines
+544
to
+547
| with pytest.raises(AssertionError): | ||
| ext.session # noqa: B018 | ||
| with pytest.raises(AssertionError): | ||
| ext.event_bus # noqa: B018 |
Collaborator
There was a problem hiding this comment.
If changes above are accepted:
Suggested change
| with pytest.raises(AssertionError): | |
| ext.session # noqa: B018 | |
| with pytest.raises(AssertionError): | |
| ext.event_bus # noqa: B018 | |
| with pytest.raises(RuntimeError): | |
| ext.session # noqa: B018 | |
| with pytest.raises(RuntimeError): | |
| ext.event_bus # noqa: B018 |
- Bound _T TypeVar to SessionExtension - Use RuntimeError instead of assert for detached property access
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This change enhances the modularity and maintainability of session extensions.