make TensorBoardWSGIApp allow passing TBPlugin or TBLoader classes directly

Author: nfelt

tensorboard/backend/application.py

@@ -160,7 +160,8 @@ def TensorBoardWSGIApp(
 
   Args:
     flags: An argparse.Namespace containing TensorBoard CLI flags.
-    plugins: A list of TBLoader subclasses for the plugins to load.
+    plugins: A list of plugins, which can be provided as TBPlugin subclasses
+        or TBLoader instances or subclasses.
     assets_zip_provider: See TBContext documentation for more information.
     data_provider: Instance of `tensorboard.data.provider.DataProvider`. May
         be `None` if `flags.generic_data` is set to `"false"` in which case
@@ -187,7 +188,8 @@ def TensorBoardWSGIApp(
       plugin_name_to_instance=plugin_name_to_instance,
       window_title=flags.window_title)
   tbplugins = []
-  for loader in plugins:
+  for plugin_spec in plugins:
+    loader = _make_plugin_loader(plugin_spec)
     plugin = loader.load(context)
     if plugin is None:
       continue
@@ -597,3 +599,21 @@ def AddRunsFromDirectory(self, path, name=None):
   def Reload(self):
     """Unsupported."""
     raise NotImplementedError()
+
+
+def _make_plugin_loader(plugin):
+  """Returns a plugin loader for the given plugin.
+
+  Args:
+    plugin: A TBPlugin subclass, or a TBLoader instance or subclass.
+
+  Returns:
+    A TBLoader for the given plugin.
+  """
+  if isinstance(plugin, base_plugin.TBLoader):
+    return plugin
+  if issubclass(plugin, base_plugin.TBLoader):
+    return plugin()
+  if issubclass(plugin, base_plugin.TBPlugin):
+    return base_plugin.BasicLoader(plugin)
+  raise ValueError("Not a TBLoader or TBPlugin subclass: %s" % plugin)

tensorboard/backend/application_test.py

@@ -135,6 +135,16 @@ def frontend_metadata(self):
     )
 
 
+class FakePluginLoader(base_plugin.TBLoader):
+  """Pass-through loader for FakePlugin with arbitrary arguments."""
+
+  def __init__(self, **kwargs):
+    self._kwargs = kwargs
+
+  def load(self, context):
+    return FakePlugin(context, **self._kwargs)
+
+
 class ApplicationTest(tb_test.TestCase):
   def setUp(self):
     plugins = [
@@ -519,23 +529,21 @@ def setUp(self):
     self.app = application.standard_tensorboard_wsgi(
         FakeFlags(logdir=self.get_temp_dir()),
         [
-          base_plugin.BasicLoader(functools.partial(
-              FakePlugin,
-              plugin_name='foo',
-              is_active_value=True,
-              routes_mapping={'/foo_route': self._foo_handler},
-              construction_callback=self._construction_callback)),
-          base_plugin.BasicLoader(functools.partial(
-              FakePlugin,
-              plugin_name='bar',
-              is_active_value=True,
-              routes_mapping={
-                  '/bar_route': self._bar_handler,
-                  '/wildcard/*': self._wildcard_handler,
-                  '/wildcard/special/*': self._wildcard_special_handler,
-                  '/wildcard/special/exact': self._foo_handler,
-              },
-              construction_callback=self._construction_callback)),
+            FakePluginLoader(
+                plugin_name='foo',
+                is_active_value=True,
+                routes_mapping={'/foo_route': self._foo_handler},
+                construction_callback=self._construction_callback),
+            FakePluginLoader(
+                plugin_name='bar',
+                is_active_value=True,
+                routes_mapping={
+                    '/bar_route': self._bar_handler,
+                    '/wildcard/*': self._wildcard_handler,
+                    '/wildcard/special/*': self._wildcard_special_handler,
+                    '/wildcard/special/exact': self._foo_handler,
+                },
+                construction_callback=self._construction_callback),
         ],
         dummy_assets_zip_provider)