From 1639b799f36d80e974a69c6ff136846c66ae6988 Mon Sep 17 00:00:00 2001
From: Joe Numainville <josephnumainville@deephaven.io>
Date: Thu, 29 Feb 2024 17:20:29 -0600
Subject: [PATCH] comments

---
 .../packaging/src/deephaven/plugin/packaging/utils.py |  2 ++
 .../plugin/utilities/dhe_safe_callback_wrapper.py     | 11 ++++++++---
 .../utilities/src/deephaven/plugin/utilities/utils.py |  3 ++-
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/plugins/packaging/src/deephaven/plugin/packaging/utils.py b/plugins/packaging/src/deephaven/plugin/packaging/utils.py
index cc4bb6f69..cfb1aa66d 100644
--- a/plugins/packaging/src/deephaven/plugin/packaging/utils.py
+++ b/plugins/packaging/src/deephaven/plugin/packaging/utils.py
@@ -2,6 +2,8 @@
 import os
 import subprocess
 
+__all__ = ["package_js"]
+
 
 def package_js(js_dir: str, dest_dir: str) -> None:
     """
diff --git a/plugins/utilities/src/deephaven/plugin/utilities/dhe_safe_callback_wrapper.py b/plugins/utilities/src/deephaven/plugin/utilities/dhe_safe_callback_wrapper.py
index d02302075..3754e2aaa 100644
--- a/plugins/utilities/src/deephaven/plugin/utilities/dhe_safe_callback_wrapper.py
+++ b/plugins/utilities/src/deephaven/plugin/utilities/dhe_safe_callback_wrapper.py
@@ -1,12 +1,17 @@
-from deephaven.plugin import Callback, Plugin
+from __future__ import annotations
+
+from typing import Type
 import logging
+from deephaven.plugin import Callback, Plugin
 from deephaven.plugin.js import JsPlugin
 from .utils import in_enterprise_environment
 
 logger = logging.getLogger(__name__)
 
+__all__ = ["DheSafeCallbackWrapper"]
+
 
-class DheSafeCallbackWrapper:
+class DheSafeCallbackWrapper(Callback):
     """
 
     A wrapper around the Callback class that provides a safe way to register plugins.
@@ -16,7 +21,7 @@ class DheSafeCallbackWrapper:
     def __init__(self, callback: Callback):
         self._callback = callback
 
-    def register(self, plugin: Plugin) -> None:
+    def register(self, plugin: Plugin | Type[Plugin]) -> None:
         """
         Register a plugin with the provided callback
 
diff --git a/plugins/utilities/src/deephaven/plugin/utilities/utils.py b/plugins/utilities/src/deephaven/plugin/utilities/utils.py
index f6d71bafb..7afce6990 100644
--- a/plugins/utilities/src/deephaven/plugin/utilities/utils.py
+++ b/plugins/utilities/src/deephaven/plugin/utilities/utils.py
@@ -1,6 +1,5 @@
 import abc
 import logging
-import os
 from functools import partial
 from typing import Callable, ContextManager
 import importlib.resources
@@ -11,6 +10,8 @@
 
 logger = logging.getLogger(__name__)
 
+__all__ = ["in_enterprise_environment", "create_js_plugin"]
+
 
 def in_enterprise_environment() -> bool:
     """