Make dicts/lists visualizable and add JSON as viz type

src/zenml/enums.py

@@ -52,6 +52,7 @@ class VisualizationType(StrEnum):
     HTML = "html"
     IMAGE = "image"
     MARKDOWN = "markdown"
+    JSON = "json"
 
 
 class ZenMLServiceType(StrEnum):

src/zenml/materializers/built_in_materializer.py

@@ -28,7 +28,7 @@
 )
 
 from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-from zenml.enums import ArtifactType
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
@@ -414,6 +414,25 @@ def save(self, data: Any) -> None:
             for entry in metadata:
                 self.artifact_store.rmtree(entry["path"])
             raise e
+
+    # save dict type objects to JSON file with JSON visualization type
+    def save_visualizations(
+        self, data: Any
+    ) -> Dict[str, "VisualizationType"]:
+        """Save visualizations for the given data.
+
+        Args:
+            data: The data to save visualizations for.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        # dict/list type objects are always saved as JSON files
+        # doesn't work for non-serializable types as they 
+        # are saved as list of lists in different files
+        if _is_serializable(data):
+            return {self.data_path: VisualizationType.JSON}
+        return {}
 
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given built-in container object.

src/zenml/materializers/structured_string_materializer.py

@@ -19,22 +19,23 @@
 from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
-from zenml.types import CSVString, HTMLString, MarkdownString
+from zenml.types import CSVString, HTMLString, JSONString, MarkdownString
 
 logger = get_logger(__name__)
 
 
-STRUCTURED_STRINGS = Union[CSVString, HTMLString, MarkdownString]
+STRUCTURED_STRINGS = Union[CSVString, HTMLString, MarkdownString, JSONString]
 
 HTML_FILENAME = "output.html"
 MARKDOWN_FILENAME = "output.md"
 CSV_FILENAME = "output.csv"
+JSON_FILENAME = "output.json"
 
 
 class StructuredStringMaterializer(BaseMaterializer):
     """Materializer for HTML or Markdown strings."""
 
-    ASSOCIATED_TYPES = (CSVString, HTMLString, MarkdownString)
+    ASSOCIATED_TYPES = (CSVString, HTMLString, MarkdownString, JSONString)
     ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA_ANALYSIS
 
     def load(self, data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS:
@@ -94,6 +95,8 @@ def _get_filepath(self, data_type: Type[STRUCTURED_STRINGS]) -> str:
             filename = HTML_FILENAME
         elif issubclass(data_type, MarkdownString):
             filename = MARKDOWN_FILENAME
+        elif issubclass(data_type, JSONString):
+            filename = JSON_FILENAME
         else:
             raise ValueError(
                 f"Data type {data_type} is not supported by this materializer."
@@ -120,6 +123,8 @@ def _get_visualization_type(
             return VisualizationType.HTML
         elif issubclass(data_type, MarkdownString):
             return VisualizationType.MARKDOWN
+        elif issubclass(data_type, JSONString):
+            return VisualizationType.JSON
         else:
             raise ValueError(
                 f"Data type {data_type} is not supported by this materializer."

src/zenml/types.py

@@ -33,3 +33,6 @@ class MarkdownString(str):
 
 class CSVString(str):
     """Special string class to indicate a CSV string."""
+
+class JSONString(str):
+    """Special string class to indicate a JSON string."""

src/zenml/utils/visualization_utils.py

@@ -13,9 +13,10 @@
 #  permissions and limitations under the License.
 """Utility functions for dashboard visualizations."""
 
+import json
 from typing import TYPE_CHECKING, Optional
 
-from IPython.core.display import HTML, Image, Markdown, display
+from IPython.core.display import HTML, Image, JSON, Markdown, display
 
 from zenml.artifacts.utils import load_artifact_visualization
 from zenml.enums import VisualizationType
@@ -63,6 +64,8 @@ def visualize_artifact(
             assert isinstance(visualization.value, str)
             table = format_csv_visualization_as_html(visualization.value)
             display(HTML(table))
+        elif visualization.type == VisualizationType.JSON:
+            display(JSON(json.loads(visualization.value)))
         else:
             display(visualization.value)