From 48213627295063135a9c003ff2f74afe5760d9c5 Mon Sep 17 00:00:00 2001 From: Ben Welsh Date: Mon, 20 Oct 2025 00:44:25 +0000 Subject: [PATCH 1/2] feat: add display method to BaseChart for rendering in Jupyter notebooks --- datawrapper/charts/base.py | 27 +++++++- tests/integration/test_base_display.py | 89 ++++++++++++++++++++++++++ 2 files changed, 115 insertions(+), 1 deletion(-) create mode 100644 tests/integration/test_base_display.py diff --git a/datawrapper/charts/base.py b/datawrapper/charts/base.py index 9a18fbc..71c64af 100644 --- a/datawrapper/charts/base.py +++ b/datawrapper/charts/base.py @@ -5,7 +5,7 @@ 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 @@ -982,6 +982,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..2a71df2 --- /dev/null +++ b/tests/integration/test_base_display.py @@ -0,0 +1,89 @@ +"""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) From 8608449e9cae969223c80aa9645ae416da817918 Mon Sep 17 00:00:00 2001 From: Ben Welsh Date: Mon, 20 Oct 2025 00:44:48 +0000 Subject: [PATCH 2/2] refactor: simplify import statements in base.py and improve code readability --- datawrapper/charts/base.py | 8 +------- tests/integration/test_base_display.py | 4 +++- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/datawrapper/charts/base.py b/datawrapper/charts/base.py index 71c64af..f57d49b 100644 --- a/datawrapper/charts/base.py +++ b/datawrapper/charts/base.py @@ -9,13 +9,7 @@ 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): diff --git a/tests/integration/test_base_display.py b/tests/integration/test_base_display.py index 2a71df2..554e174 100644 --- a/tests/integration/test_base_display.py +++ b/tests/integration/test_base_display.py @@ -69,7 +69,9 @@ def test_base_chart_display_with_access_token(): 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: + 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)