Merge pull request #461 from palewire/main

feat: add display method to BaseChart for rendering in Jupyter notebooks

Author: palewire

datawrapper/charts/base.py

@@ -5,17 +5,11 @@
 from typing import Any, Literal
 
 import pandas as pd
-from IPython.display import Image
+from IPython.display import IFrame, Image
 from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator
 
 from datawrapper.__main__ import Datawrapper
-from datawrapper.charts.models import (
-    Annotate,
-    Describe,
-    Publish,
-    Transform,
-    Visualize,
-)
+from datawrapper.charts.models import Annotate, Describe, Publish, Transform, Visualize
 
 
 class BaseChart(BaseModel):
@@ -982,6 +976,31 @@ def get_iframe_code(
         # Call the get_iframe_code method from the client
         return client.get_iframe_code(chart_id=self.chart_id, responsive=responsive)
 
+    def display(self, access_token: str | None = None) -> IFrame:
+        """Display the chart as an IFrame in a Jupyter notebook.
+
+        Args:
+            access_token: Optional Datawrapper API access token.
+                         If not provided, will use DATAWRAPPER_ACCESS_TOKEN environment variable.
+
+        Returns:
+            An IPython.display.IFrame object displaying the chart.
+
+        Raises:
+            ValueError: If no chart_id is set or no access token is available.
+            Exception: If the API request fails.
+        """
+        if not self.chart_id:
+            raise ValueError(
+                "No chart_id set. Use create() first or set chart_id manually."
+            )
+
+        # Get the client
+        client = self._get_client(access_token)
+
+        # Call the display_chart method from the client
+        return client.display_chart(chart_id=self.chart_id)
+
     def get_editor_url(self) -> str:
         """Get the Datawrapper editor URL for this chart.
 

tests/integration/test_base_display.py

@@ -0,0 +1,91 @@
+"""Test the display method on BaseChart."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from IPython.display import IFrame
+
+from datawrapper.charts.base import BaseChart
+
+
+def test_base_chart_display_method_exists():
+    """Test that the display method exists on BaseChart."""
+    chart = BaseChart(chart_type="d3-lines", title="Test Chart")
+    assert hasattr(chart, "display")
+    assert callable(chart.display)
+
+
+def test_base_chart_display_requires_chart_id():
+    """Test that display raises ValueError when chart_id is not set."""
+    chart = BaseChart(chart_type="d3-lines", title="Test Chart")
+
+    with pytest.raises(ValueError, match="No chart_id set"):
+        chart.display()
+
+
+def test_base_chart_display_with_chart_id():
+    """Test that display works when chart_id is set."""
+    # Create a chart with a chart_id
+    chart = BaseChart(
+        chart_type="d3-lines",
+        title="Test Display Chart",
+        data=[{"x": 1, "y": 2}, {"x": 2, "y": 4}],
+    )
+    chart.chart_id = "test123"
+
+    # Mock the client and its display_chart method
+    mock_client = MagicMock()
+    mock_iframe = IFrame("https://example.com", width=600, height=400)
+    mock_client.display_chart.return_value = mock_iframe
+
+    with patch.object(chart, "_get_client", return_value=mock_client):
+        # Call the display method
+        result = chart.display()
+
+        # Verify the client method was called
+        mock_client.display_chart.assert_called_once()
+
+        # Verify the chart_id was passed
+        call_kwargs = mock_client.display_chart.call_args.kwargs
+        assert call_kwargs["chart_id"] == "test123"
+
+        # Verify the result is an IFrame
+        assert result == mock_iframe
+        assert isinstance(result, IFrame)
+
+
+def test_base_chart_display_with_access_token():
+    """Test that display passes access_token to the client."""
+    # Create a chart with a chart_id
+    chart = BaseChart(
+        chart_type="d3-lines",
+        title="Test Display Chart with Token",
+        data=[{"x": 1, "y": 2}, {"x": 2, "y": 4}],
+    )
+    chart.chart_id = "test123"
+
+    # Mock the client and its display_chart method
+    mock_client = MagicMock()
+    mock_iframe = IFrame("https://example.com", width=600, height=400)
+    mock_client.display_chart.return_value = mock_iframe
+
+    with patch.object(
+        chart, "_get_client", return_value=mock_client
+    ) as mock_get_client:
+        # Call the display method with an access token
+        test_token = "test_access_token"
+        result = chart.display(access_token=test_token)
+
+        # Verify _get_client was called with the access token
+        mock_get_client.assert_called_once_with(test_token)
+
+        # Verify the client method was called
+        mock_client.display_chart.assert_called_once()
+
+        # Verify the chart_id was passed
+        call_kwargs = mock_client.display_chart.call_args.kwargs
+        assert call_kwargs["chart_id"] == "test123"
+
+        # Verify the result is an IFrame
+        assert result == mock_iframe
+        assert isinstance(result, IFrame)