Merge pull request #591 from JohnBarton27/feature/add-pie-charts

Add initial Pie Chart implementation

pygsheets/custom_types.py

@@ -123,3 +123,4 @@ class ChartType(Enum):
     SCATTER = "SCATTER"
     COMBO = "COMBO"
     STEPPED_AREA = "STEPPED_AREA"
+    PIE = "PIE"

pygsheets/pie_chart.py

@@ -0,0 +1,111 @@
+from pygsheets.chart import Chart
+
+
+class PieChart(Chart):
+    """
+    Represents a Pie Chart in a worksheet.
+
+    :param worksheet:         Worksheet object in which the chart resides
+    :param domain:            Cell range of the desired chart domain in the form of tuple of tuples
+    :param chart_range:       Cell ranges of the desired (singular) range in the form of a tuple of tuples
+    :param title:             Title of the chart
+    :param anchor_cell:       Position of the left corner of the chart in the form of cell address or cell object
+    :param three_dimensional  True if the pie is three dimensional
+    :param pie_hole           (float) The size of the hole in the pie chart (defaults to 0). Must be between 0 and 1.
+    :param json_obj:          Represents a json structure of the chart as given in `api <https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets#BasicChartSpec>`__.
+    """
+    def __init__(self, worksheet, domain=None, chart_range=None, title='', anchor_cell=None, three_dimensional=False,
+                 pie_hole=0, json_obj=None):
+        self._three_dimensional = three_dimensional
+        self._pie_hole = pie_hole
+
+        if self._pie_hole < 0 or self._pie_hole > 1:
+            raise ValueError("Pie Chart's pie_hole must be between 0 and 1.")
+
+        super().__init__(worksheet, domain, ranges=[chart_range], chart_type=None, title=title,
+                         anchor_cell=anchor_cell, json_obj=json_obj)
+
+    def get_json(self):
+        """Returns the pie chart as a dictionary structured like the Google Sheets API v4."""
+
+        domains = [{'domain': {'sourceRange': {'sources': [
+            self._worksheet.get_gridrange(self._domain[0], self._domain[1])]}}}]
+        ranges = self._get_ranges_request()
+        spec = dict()
+        spec['title'] = self._title
+        spec['pieChart'] = dict()
+        spec['pieChart']['legendPosition'] = self._legend_position
+        spec['fontName'] = self._font_name
+        spec['pieChart']['domain'] = domains[0]["domain"]
+        spec['pieChart']['series'] = ranges[0]["series"]
+        spec['pieChart']['threeDimensional'] = self._three_dimensional
+        spec['pieChart']['pieHole'] = self._pie_hole
+        return spec
+
+    def _create_chart(self):
+        domains = []
+        if self._domain:
+            domains.append({
+                "domain": {
+                    "sourceRange": {
+                        "sources": [self._worksheet.get_gridrange(self._domain[0], self._domain[1])]
+                    }
+                }
+            })
+
+        request = {
+            "addChart": {
+                "chart": {
+                    "spec": {
+                        "title": self._title,
+                        "pieChart": {
+                            "domain": domains[0]["domain"] if domains else None,
+                            "series": self._get_ranges_request()[0]["series"],
+                            "threeDimensional": self._three_dimensional,
+                            "pieHole": self._pie_hole
+                        },
+                    },
+                    "position": {
+                        "overlayPosition": {
+                            "anchorCell": self._get_anchor_cell()
+                        }
+                    }
+                }
+            }
+        }
+        response = self._worksheet.client.sheet.batch_update(self._worksheet.spreadsheet.id, request)
+        chart_data_list = response.get('replies')
+        chart_json = chart_data_list[0].get('addChart',{}).get('chart')
+        self.set_json(chart_json)
+
+    def set_json(self, chart_data):
+        """
+        Reads a json-dictionary returned by the Google Sheets API v4 and initialize all the properties from it.
+
+        :param chart_data:   The chart data as json specified in sheets api.
+        """
+        anchor_cell_data = chart_data.get('position',{}).get('overlayPosition',{}).get('anchorCell')
+        self._anchor_cell = (anchor_cell_data.get('rowIndex',0)+1, anchor_cell_data.get('columnIndex',0)+1)
+        self._title = chart_data.get('spec',{}).get('title',None)
+        self._chart_id = chart_data.get('chartId',None)
+        self._title_font_family = chart_data.get('spec',{}).get('titleTextFormat',{}).get('fontFamily',None)
+        self._font_name = chart_data.get('spec',{}).get('titleTextFormat',{}).get('fontFamily',None)
+        pie_chart = chart_data.get('spec',{}).get('pieChart', None)
+        self._legend_position = pie_chart.get('legendPosition', None)
+        domain = pie_chart.get('domain', {})
+        source_list = domain.get('sourceRange', {}).get('sources', None)
+        for source in source_list:
+            start_row = source.get('startRowIndex',0)
+            end_row = source.get('endRowIndex',0)
+            start_column = source.get('startColumnIndex',0)
+            end_column = source.get('endColumnIndex',0)
+            self._domain = [(start_row+1, start_column+1),(end_row, end_column)]
+        range = pie_chart.get('series', {})
+        self._ranges = []
+        source_list = range.get('sourceRange',{}).get('sources',None)
+        for source in source_list:
+            start_row = source.get('startRowIndex',0)
+            end_row = source.get('endRowIndex',0)
+            start_column = source.get('startColumnIndex',0)
+            end_column = source.get('endColumnIndex',0)
+            self._ranges.append([(start_row+1, start_column+1), (end_row, end_column)])

pygsheets/worksheet.py

@@ -20,6 +20,7 @@
 from pygsheets.utils import numericise_all, format_addr, fullmatch, batchable, allow_gridrange, get_color_style, get_boolean_condition
 from pygsheets.custom_types import *
 from pygsheets.chart import Chart
+from pygsheets.pie_chart import PieChart
 from pygsheets.developer_metadata import DeveloperMetadataLookupDataFilter, DeveloperMetadata
 try:
     import pandas as pd
@@ -1657,9 +1658,9 @@ def add_chart(self, domain, ranges, title=None, chart_type=ChartType.COLUMN, anc
         You can just add the rainfall data as a range.
 
 
-        :param domain:          Cell range of the desired chart domain (x-axis) in the form of tuple of adresses
+        :param domain:          Cell range of the desired chart domain (x-axis) in the form of tuple of addresses
                                 (start_address, end_address)
-        :param ranges:          Cell ranges of the desired ranges (y-axis) in the form of list of tuples of adresses
+        :param ranges:          Cell ranges of the desired ranges (y-axis) in the form of list of tuples of addresses
         :param title:           Title of the chart
         :param chart_type:      Basic chart type (default: COLUMN)
         :param anchor_cell:     position of the left corner of the chart in the form of cell address or cell object
@@ -1676,6 +1677,19 @@ def add_chart(self, domain, ranges, title=None, chart_type=ChartType.COLUMN, anc
         """
         return Chart(self, domain, ranges, chart_type, title, anchor_cell)
 
+    def add_pie_chart(self, domain, chart_range, title=None, anchor_cell=None, three_dimensional=False, pie_hole=0):
+        """
+        Similar to `add_chart`, but created a Pie Chart instead of a Basic Chart.
+        :param domain: Cell range of the desired chart domain (x-axis) in the form of tuple of addresses (start_address, end_address)
+        :param chart_range: Cell ranges of the desired (singular) range (y-axis) in the form of tuples of addresses
+        :param title: Title of the chart
+        :param anchor_cell: position of the left corner of the chart in the form of cell address or cell object
+        :param three_dimensional: True if the pie is three dimensional
+        :param pie_hole: (float) The size of the hole in the pie chart (defaults to 0). Must be between 0 and 1.
+        :return: :class:`PieChart`
+        """
+        return PieChart(self, domain, chart_range, title, anchor_cell, three_dimensional, pie_hole)
+
     def get_charts(self, title=None):
         """Returns a list of chart objects, can be filtered by title.
 

tests/online_test.py

@@ -731,6 +731,64 @@ def test_add_chart(self):
         obj.delete()
         self.worksheet.clear()
 
+    def test_add_pie_chart(self):
+        self.worksheet.resize(50,50)
+        self.worksheet.update_values('A10:C13', [['x', 'y', 'z'], [1, 5, 9]])
+        dmn = [(10, 1), (13, 1)]
+        rng = [(10, 2), (13, 2)]
+        obj = self.worksheet.add_pie_chart(dmn, rng, "Test Pie Chart", "A16")
+        assert obj.title == "Test Pie Chart"
+        assert obj.domain == dmn
+        assert obj.ranges[0] == rng
+        assert obj._three_dimensional is False
+        assert obj._pie_hole == 0
+        assert obj.font_name == "Roboto"
+        assert obj.title_font_family == "Roboto"
+        obj.delete()
+        self.worksheet.clear()
+
+    def test_add_pie_chart_three_dimensional(self):
+        self.worksheet.resize(50,50)
+        self.worksheet.update_values('A10:C13', [['x', 'y', 'z'], [1, 5, 9]])
+        dmn = [(10, 1), (13, 1)]
+        rng = [(10, 2), (13, 2)]
+        obj = self.worksheet.add_pie_chart(dmn, rng, "Test Pie Chart", "A16", three_dimensional=True)
+        assert obj.title == "Test Pie Chart"
+        assert obj.domain == dmn
+        assert obj.ranges[0] == rng
+        assert obj._three_dimensional is True
+        assert obj._pie_hole == 0
+        assert obj.font_name == "Roboto"
+        assert obj.title_font_family == "Roboto"
+        obj.delete()
+        self.worksheet.clear()
+
+    def test_add_pie_chart_pie_hole(self):
+        self.worksheet.resize(50,50)
+        self.worksheet.update_values('A10:C13', [['x', 'y', 'z'], [1, 5, 9]])
+        dmn = [(10, 1), (13, 1)]
+        rng = [(10, 2), (13, 2)]
+        obj = self.worksheet.add_pie_chart(dmn, rng, "Test Pie Chart", "A16", pie_hole=0.5)
+        assert obj.title == "Test Pie Chart"
+        assert obj.domain == dmn
+        assert obj.ranges[0] == rng
+        assert obj._three_dimensional is False
+        assert obj._pie_hole == 0.5
+        assert obj.font_name == "Roboto"
+        assert obj.title_font_family == "Roboto"
+        obj.delete()
+        self.worksheet.clear()
+
+    def test_add_pie_chart_invalid_pie_hole(self):
+        self.worksheet.resize(50,50)
+        self.worksheet.update_values('A10:C13', [['x', 'y', 'z'], [1, 5, 9]])
+        dmn = [(10, 1), (13, 1)]
+        rng = [(10, 2), (13, 2)]
+        with pytest.raises(ValueError):
+            obj = self.worksheet.add_pie_chart(dmn, rng, "Test Pie Chart", "A16", pie_hole=2)
+
+        self.worksheet.clear()
+
     def test_get_charts(self):
         self.worksheet.resize(50,50)
         self.worksheet.update_values('A30:C33',[['x','y','z'],[1,5,9],[2,4,8],[3,6,10]])