diff --git a/datawrapper/__main__.py b/datawrapper/__main__.py index bcf3c12..2d9aa98 100644 --- a/datawrapper/__main__.py +++ b/datawrapper/__main__.py @@ -16,8 +16,9 @@ import json import logging import os +from io import StringIO from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Union +from typing import Any, Iterable import IPython import pandas as pd @@ -58,7 +59,7 @@ def __init__(self, access_token=_ACCESS_TOKEN): self._access_token = access_token self._auth_header = {"Authorization": f"Bearer {access_token}"} - def account_info(self) -> Union[Dict[Any, Any], None, Any]: + def account_info(self) -> dict[Any, Any] | None | Any: """Access your account information. Returns @@ -138,8 +139,8 @@ def create_chart( data: pd.DataFrame | str | None = None, folder_id: str = "", organization_id: str = "", - metadata: Optional[Dict[Any, Any]] = None, - ) -> Union[Dict[Any, Any], None, Any]: + metadata: dict[Any, Any] | None = None, + ) -> dict[Any, Any] | None | Any: """Creates a new Datawrapper chart, table or map. You can pass a pandas DataFrame as a `data` argument to upload data. @@ -220,7 +221,7 @@ def update_description( number_append: str = "", number_format: str = "-", number_divisor: int = 0, - ) -> Union[Any, None]: + ) -> Any | None: """Update a chart's description. Parameters @@ -278,7 +279,7 @@ def update_description( logger.error("Couldn't update chart.") return None - def publish_chart(self, chart_id: str, display: bool = True) -> Union[Any, None]: + def publish_chart(self, chart_id: str, display: bool = True) -> Any | None: """Publishes a chart, table or map. Parameters @@ -311,7 +312,7 @@ def publish_chart(self, chart_id: str, display: bool = True) -> Union[Any, None] def chart_properties( self, chart_id: str - ) -> Union[Dict[Any, Any], None, Any, Iterable[Any]]: + ) -> dict[Any, Any] | None | Any | Iterable[Any]: """Retrieve information of a specific chart, table or map. Parameters @@ -356,14 +357,21 @@ def chart_data(self, chart_id: str): ) # Check if the request was successful - assert response.ok, "Make sure you have the right id and authorization credentials (access_token)." - - # Return the data - return response.json() + assert ( + response.ok + ), "Make sure you have the right id and authorization credentials (access_token)." + + # Return the data as json if the mimetype is json + if "json" in response.headers["content-type"]: + return response.json() + # If it's a csv, read the text into a dataframe + elif "text/csv" in response.headers["content-type"]: + return pd.read_csv(StringIO(response.text)) + # Otherwise just return the text + else: + return response.text - def update_metadata( - self, chart_id: str, properties: Dict[Any, Any] - ) -> Union[Any, None]: + def update_metadata(self, chart_id: str, properties: dict[Any, Any]) -> Any | None: """Update a chart, table, or map's metadata. Example: https://developer.datawrapper.de/docs/creating-a-chart-new#edit-colors @@ -404,7 +412,7 @@ def update_chart( language: str = "", folder_id: str = "", organization_id: str = "", - ) -> Union[Any, None]: + ) -> Any | None: """Updates a chart's title, theme, type, language, or location (folder/organization). Parameters @@ -473,9 +481,7 @@ def display_chart(self, chart_id: str) -> IPython.display.HTML: return HTML(_iframe_code) - def get_iframe_code( - self, chart_id: str, responsive: bool = False - ) -> Union[str, Any]: + def get_iframe_code(self, chart_id: str, responsive: bool = False) -> str | Any: """Returns a chart, table, or map's iframe embed code. Parameters @@ -516,7 +522,7 @@ def export_chart( output: str = "png", filepath: str = "./image.png", display: bool = False, - ) -> Union[Any, None]: + ) -> Any | None: """Exports a chart, table, or map. Parameters @@ -564,7 +570,7 @@ def export_chart( "zoom": zoom, "scale": scale, "borderWidth": border_width, - "transparent": transparent + "transparent": transparent, } _header = self._auth_header @@ -595,7 +601,7 @@ def export_chart( logger.error(msg) raise Exception(msg) - def get_folders(self) -> Union[Dict[Any, Any], None, Any]: + def get_folders(self) -> dict[Any, Any] | None | Any: """Get a list of folders in your Datawrapper account. Returns @@ -616,7 +622,7 @@ def get_folders(self) -> Union[Dict[Any, Any], None, Any]: ) return None - def move_chart(self, chart_id: str, folder_id: str) -> Union[Any, None]: + def move_chart(self, chart_id: str, folder_id: str) -> Any | None: """Moves a chart, table, or map to a specified folder. Parameters @@ -677,7 +683,7 @@ def get_charts( limit: int = 25, folder_id: str = "", team_id: str = "", - ) -> Union[None, List[Any]]: + ) -> None | list[Any]: """Retrieves a list of charts by User Parameters