Introduce an add_json method for uploading markers, areas and lines t…

…o locator maps
chekos · Nov 8, 2023 · e17c5c6 · e17c5c6
1 parent fb72960
commit e17c5c6
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 25 deletions.
diff --git a/datawrapper/__main__.py b/datawrapper/__main__.py
@@ -1106,6 +1106,37 @@ def add_data(self, chart_id: str, data: pd.DataFrame | str) -> bool:
             dump_data=False,
         )
 
+    def add_json(self, chart_id: str, data: Any) -> bool:
+        """Add JSON data to a specified chart.
+
+        Can be used to add point, area and line markers to a locator map or other chart.
+
+        Parameters
+        ----------
+        chart_id : str
+            ID of chart, table or map to add data to.
+        data : Any
+            JSON data to add to the chart.
+
+        Returns
+        -------
+        bool
+            True if the data was added successfully.
+        """
+        # Set the chart metadata to accept JSON data
+        self.update_chart(
+            chart_id=chart_id,
+            metadata={
+                "data": {"json": True},
+            },
+        )
+
+        # Dump the provided data as a JSON string
+        json_data = json.dumps(data)
+
+        # Post it to the chart via the add_data method
+        return self.add_data(chart_id, json_data)
+
     def refresh_data(self, chart_id: str) -> dict:
         """Fetch configured external data and add it to the chart.
 

diff --git a/tests/test_maps.py b/tests/test_maps.py
@@ -54,7 +54,7 @@ def test_locator_map_points():
     dw.update_chart(obj["id"], metadata=metadata)
 
     # Add markers
-    point_markers = {
+    data = {
         "markers": [
             {
                 "type": "point",
@@ -80,8 +80,9 @@ def test_locator_map_points():
             }
         ]
     }
-    data = json.dumps(point_markers)
-    dw.add_data(obj["id"], data)
+
+    # Post it
+    dw.add_json(obj["id"], data)
 
 
 def test_locator_map_areas():
@@ -95,19 +96,12 @@ def test_locator_map_areas():
         "locator-map",
     )
 
-    # Set map metadata
-    dw.update_chart(
-        obj["id"],
-        metadata={
-            "data": {"json": True},
-        },
-    )
+    # Open markers
+    with open("./tests/area_markers.json") as f:
+        data = json.load(f)
 
     # Add markers
-    with open("./tests/area_markers.json") as f:
-        area_markers = json.load(f)
-    data = json.dumps(area_markers)
-    dw.add_data(obj["id"], data)
+    dw.add_json(obj["id"], data)
 
 
 def test_locator_map_lines():
@@ -121,16 +115,8 @@ def test_locator_map_lines():
         "locator-map",
     )
 
-    # Set map metadata
-    dw.update_chart(
-        obj["id"],
-        metadata={
-            "data": {"json": True},
-        },
-    )
-
     # Add markers
-    line_markers = {
+    data = {
         "markers": [
             {
                 "id": "m1",
@@ -155,5 +141,4 @@ def test_locator_map_lines():
             }
         ]
     }
-    data = json.dumps(line_markers)
-    dw.add_data(obj["id"], data)
+    dw.add_json(obj["id"], data)