def testVisualizationRegistrationAndUsage(self, *unused_mocks):
registry = visualizations.ArtifactVisualizationRegistry()
visualizations.get_registry = mock.MagicMock(return_value=registry)
mock_object = mock.MagicMock()
class MyVisualization(visualizations.ArtifactVisualization):
# Arbitrary artifact type class.
ARTIFACT_TYPE = standard_artifacts.Examples
def display(self, unused_artifact):
mock_object('foo')
visualizations.get_registry().register(MyVisualization)
self.assertIs(
MyVisualization,
visualizations.get_registry().get_visualization(
standard_artifacts.Examples.TYPE_NAME).__class__)
def show(self, item: object) -> None:
"""Show the given object in an IPython notebook display."""
from IPython.core.display import display # pylint: disable=g-import-not-at-top
from IPython.core.display import HTML # pylint: disable=g-import-not-at-top
if isinstance(item, types.Channel):
channel = item
artifacts = channel.get()
for artifact in artifacts:
artifact_heading = 'Artifact at %s' % cgi.escape(artifact.uri)
display(HTML('<b>%s</b><br/><br/>' % artifact_heading))
visualization = visualizations.get_registry(
).get_visualization(artifact.type_name)
if visualization:
visualization.display(artifact)
else:
display(item)
def register_standard_visualizations():
for visualization in STANDARD_VISUALIZATIONS:
visualizations.get_registry().register(visualization)
