diff --git a/datawrapper/charts/base.py b/datawrapper/charts/base.py index 9a18fbc..f57d49b 100644 --- a/datawrapper/charts/base.py +++ b/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. diff --git a/tests/integration/test_base_display.py b/tests/integration/test_base_display.py new file mode 100644 index 0000000..554e174 --- /dev/null +++ b/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)