Repository: plotly/plotly.py
Pull request #588 (2016-10-21 12:12:25)

Gantt chart grouping
Added an option to group Gantt Charts tasks.
Also added an option to add Texts to tasks in the chart via a 'Description' key.

The tests all ran successfully and I added another one for the new options: `test_gannt_groups_and_descriptions`
Test Location: **plotly/tests/test_core/test_tools/test_figure_factory.py**


```python
df = [
    dict(Task='Task A', Description='Task A - 1', Start='2008-10-05',
         Finish='2009-04-15', IndexCol = 'TA'),
    dict(Task="Task B", Description='Task B - 1', Start='2008-12-06',
         Finish='2009-03-15', IndexCol = 'TB'),
    dict(Task="Task C", Description='Task C - 1', Start='2008-09-07',
         Finish='2009-03-15', IndexCol = 'TC'),
    dict(Task="Task C", Description='Task C - 2', Start='2009-05-08',
         Finish='2009-04-15', IndexCol = 'TC'),
    dict(Task="Task A", Description='Task A - 2', Start='2009-04-20',
         Finish='2009-05-30', IndexCol = 'TA')
]
fig = FigureFactory.create_gantt(df, colors=dict(TA = 'rgb(220, 0, 0)',TB = 'rgb(170, 14, 200)',TC = (1, 0.9, 0.16)), show_colorbar=True, index_col='IndexCol', group_tasks=True)
```

Generates:

![gantt-chart-grouping](https://cloud.githubusercontent.com/assets/1822673/19597826/4a740aac-9797-11e6-8f67-fd92b36fc2c6.png)

Did I forget anything?


Patch:
diff --git a/plotly/tools.py b/plotly/tools.py
index e8af25dc5..30d8dba02 100644
--- a/plotly/tools.py
+++ b/plotly/tools.py
@@ -1724,7 +1724,7 @@ class FigureFactory(object):
 
     @staticmethod
     def _gantt(chart, colors, title, bar_width, showgrid_x, showgrid_y,
-               height, width, tasks=None, task_names=None, data=None):
+               height, width, tasks=None, task_names=None, data=None, group_tasks=False):
         """
         Refer to FigureFactory.create_gantt() for docstring
         """
@@ -1739,6 +1739,8 @@ class FigureFactory(object):
             task = dict(x0=chart[index]['Start'],
                         x1=chart[index]['Finish'],
                         name=chart[index]['Task'])
+            if 'Description' in chart[index]:
+                task['description'] = chart[index]['Description']
             tasks.append(task)
 
         shape_template = {
@@ -1751,29 +1753,49 @@ class FigureFactory(object):
             },
             'yref': 'y',
         }
+        # create the list of task names
+        for index in range(len(tasks)):
+            tn = tasks[index]['name']
+            # Is added to task_names if group_tasks is set to False,
+            # or if the option is used (True) it only adds them if the
+            # name is not already in the list
+            if not group_tasks or tn not in task_names:
+                task_names.append(tn)
+        # Guarantees that for grouped tasks the tasks that are inserted first
+        # are shown at the top
+        if group_tasks:
+            task_names.reverse()
+
 
         color_index = 0
         for index in range(len(tasks)):
             tn = tasks[index]['name']
-            task_names.append(tn)
             del tasks[index]['name']
             tasks[index].update(shape_template)
-            tasks[index]['y0'] = index - bar_width
-            tasks[index]['y1'] = index + bar_width
+
+            # If group_tasks is True, all tasks with the same name belong
+            # to the same row.
+            groupID = index
+            if group_tasks:
+                groupID = task_names.index(tn)
+            tasks[index]['y0'] = groupID - bar_width
+            tasks[index]['y1'] = groupID + bar_width
 
             # check if colors need to be looped
             if color_index >= len(colors):
                 color_index = 0
             tasks[index]['fillcolor'] = colors[color_index]
             # Add a line for hover text and autorange
-            data.append(
-                dict(
+            entry = dict(
                     x=[tasks[index]['x0'], tasks[index]['x1']],
-                    y=[index, index],
+                    y=[groupID, groupID],
                     name='',
                     marker={'color': 'white'}
                 )
-            )
+            if "description" in tasks[index]:
+                entry['text'] = tasks[index]['description']
+                del tasks[index]['description']
+            data.append(entry)
             color_index += 1
 
         layout = dict(
@@ -1786,8 +1808,8 @@ class FigureFactory(object):
             yaxis=dict(
                 showgrid=showgrid_y,
                 ticktext=task_names,
-                tickvals=list(range(len(tasks))),
-                range=[-1, len(tasks) + 1],
+                tickvals=list(range(len(task_names))),
+                range=[-1, len(task_names) + 1],
                 autorange=False,
                 zeroline=False,
             ),
@@ -1830,7 +1852,7 @@ class FigureFactory(object):
     @staticmethod
     def _gantt_colorscale(chart, colors, title, index_col, show_colorbar,
                           bar_width, showgrid_x, showgrid_y, height,
-                          width, tasks=None, task_names=None, data=None):
+                          width, tasks=None, task_names=None, data=None, group_tasks=False):
         """
         Refer to FigureFactory.create_gantt() for docstring
         """
@@ -1847,6 +1869,8 @@ class FigureFactory(object):
             task = dict(x0=chart[index]['Start'],
                         x1=chart[index]['Finish'],
                         name=chart[index]['Task'])
+            if 'Description' in chart[index]:
+                task['description'] = chart[index]['Description']
             tasks.append(task)
 
         shape_template = {
@@ -1870,13 +1894,32 @@ class FigureFactory(object):
                     "colors given will be used for the lower and upper "
                     "bounds on the colormap."
                 )
+
+            # create the list of task names
+            for index in range(len(tasks)):
+                tn = tasks[index]['name']
+                # Is added to task_names if group_tasks is set to False,
+                # or if the option is used (True) it only adds them if the
+                # name is not already in the list
+                if not group_tasks or tn not in task_names:
+                    task_names.append(tn)
+            # Guarantees that for grouped tasks the tasks that are inserted
+            # first are shown at the top
+            if group_tasks:
+                task_names.reverse()
+
             for index in range(len(tasks)):
                 tn = tasks[index]['name']
-                task_names.append(tn)
                 del tasks[index]['name']
                 tasks[index].update(shape_template)
-                tasks[index]['y0'] = index - bar_width
-                tasks[index]['y1'] = index + bar_width
+
+                # If group_tasks is True, all tasks with the same name belong
+                # to the same row.
+                groupID = index
+                if group_tasks:
+                    groupID = task_names.index(tn)
+                tasks[index]['y0'] = groupID - bar_width
+                tasks[index]['y1'] = groupID + bar_width
 
                 # unlabel color
                 colors = FigureFactory._color_parser(
@@ -1899,14 +1942,17 @@ class FigureFactory(object):
                 )
 
                 # add a line for hover text and autorange
-                data.append(
-                    dict(
+                entry = dict(
                         x=[tasks[index]['x0'], tasks[index]['x1']],
-                        y=[index, index],
+                        y=[groupID, groupID],
                         name='',
                         marker={'color': 'white'}
                     )
-                )
+                if "description" in tasks[index]:
+                    entry['text'] = tasks[index]['description']
+                    del tasks[index]['description']
+                data.append(entry)
+
 
             if show_colorbar is True:
             # generate dummy data for colorscale visibility
@@ -1948,27 +1994,46 @@ class FigureFactory(object):
                 index_vals_dict[key] = colors[c_index]
                 c_index += 1
 
+            # create the list of task names
+            for index in range(len(tasks)):
+                tn = tasks[index]['name']
+                # Is added to task_names if group_tasks is set to False,
+                # or if the option is used (True) it only adds them if the
+                # name is not already in the list
+                if not group_tasks or tn not in task_names:
+                    task_names.append(tn)
+            # Guarantees that for grouped tasks the tasks that are inserted
+            # first are shown at the top
+            if group_tasks:
+                task_names.reverse()
+
             for index in range(len(tasks)):
                 tn = tasks[index]['name']
-                task_names.append(tn)
                 del tasks[index]['name']
                 tasks[index].update(shape_template)
-                tasks[index]['y0'] = index - bar_width
-                tasks[index]['y1'] = index + bar_width
+                # If group_tasks is True, all tasks with the same name belong
+                # to the same row.
+                groupID = index
+                if group_tasks:
+                    groupID = task_names.index(tn)
+                tasks[index]['y0'] = groupID - bar_width
+                tasks[index]['y1'] = groupID + bar_width
 
                 tasks[index]['fillcolor'] = index_vals_dict[
                     chart[index][index_col]
                 ]
 
                 # add a line for hover text and autorange
-                data.append(
-                    dict(
+                entry = dict(
                         x=[tasks[index]['x0'], tasks[index]['x1']],
-                        y=[index, index],
+                        y=[groupID, groupID],
                         name='',
                         marker={'color': 'white'}
                     )
-                )
+                if "description" in tasks[index]:
+                    entry['text'] = tasks[index]['description']
+                    del tasks[index]['description']
+                data.append(entry)
 
             if show_colorbar is True:
             # generate dummy data to generate legend
@@ -1998,8 +2063,8 @@ class FigureFactory(object):
             yaxis=dict(
                 showgrid=showgrid_y,
                 ticktext=task_names,
-                tickvals=list(range(len(tasks))),
-                range=[-1, len(tasks) + 1],
+                tickvals=list(range(len(task_names))),
+                range=[-1, len(task_names) + 1],
                 autorange=False,
                 zeroline=False,
             ),
@@ -2042,7 +2107,7 @@ class FigureFactory(object):
     @staticmethod
     def _gantt_dict(chart, colors, title, index_col, show_colorbar, bar_width,
                     showgrid_x, showgrid_y, height, width, tasks=None,
-                    task_names=None, data=None):
+                    task_names=None, data=None, group_tasks=False):
         """
         Refer to FigureFactory.create_gantt() for docstring
         """
@@ -2058,6 +2123,8 @@ class FigureFactory(object):
             task = dict(x0=chart[index]['Start'],
                         x1=chart[index]['Finish'],
                         name=chart[index]['Task'])
+            if 'Description' in chart[index]:
+                task['description'] = chart[index]['Description']
             tasks.append(task)
 
         shape_template = {
@@ -2086,25 +2153,45 @@ class FigureFactory(object):
                     "keys must be all the values in the index column."
                 )
 
+        # create the list of task names
+        for index in range(len(tasks)):
+            tn = tasks[index]['name']
+            # Is added to task_names if group_tasks is set to False,
+            # or if the option is used (True) it only adds them if the
+            # name is not already in the list
+            if not group_tasks or tn not in task_names:
+                task_names.append(tn)
+        # Guarantees that for grouped tasks the tasks that are inserted first
+        # are shown at the top
+        if group_tasks:
+            task_names.reverse()
+
         for index in range(len(tasks)):
             tn = tasks[index]['name']
-            task_names.append(tn)
             del tasks[index]['name']
             tasks[index].update(shape_template)
-            tasks[index]['y0'] = index - bar_width
-            tasks[index]['y1'] = index + bar_width
+
+            # If group_tasks is True, all tasks with the same name belong
+            # to the same row.
+            groupID = index
+            if group_tasks:
+                groupID = task_names.index(tn)
+            tasks[index]['y0'] = groupID - bar_width
+            tasks[index]['y1'] = groupID + bar_width
 
             tasks[index]['fillcolor'] = colors[chart[index][index_col]]
 
             # add a line for hover text and autorange
-            data.append(
-                dict(
+            entry = dict(
                     x=[tasks[index]['x0'], tasks[index]['x1']],
-                    y=[index, index],
+                    y=[groupID, groupID],
                     name='',
                     marker={'color': 'white'}
                 )
-            )
+            if "description" in tasks[index]:
+                entry['text'] = tasks[index]['description']
+                del tasks[index]['description']
+            data.append(entry)
 
         if show_colorbar is True:
         # generate dummy data to generate legend
@@ -2134,8 +2221,8 @@ class FigureFactory(object):
             yaxis=dict(
                 showgrid=showgrid_y,
                 ticktext=task_names,
-                tickvals=list(range(len(tasks))),
-                range=[-1, len(tasks) + 1],
+                tickvals=list(range(len(task_names))),
+                range=[-1, len(task_names) + 1],
                 autorange=False,
                 zeroline=False,
             ),
@@ -2180,7 +2267,7 @@ class FigureFactory(object):
                      reverse_colors=False, title='Gantt Chart',
                      bar_width=0.2, showgrid_x=False, showgrid_y=False,
                      height=600, width=900, tasks=None,
-                     task_names=None, data=None):
+                     task_names=None, data=None, group_tasks=False):
         """
         Returns figure for a gantt chart
 
@@ -2359,7 +2446,8 @@ class FigureFactory(object):
                 )
             fig = FigureFactory._gantt(
                 chart, colors, title, bar_width, showgrid_x, showgrid_y,
-                height, width, tasks=None, task_names=None, data=None
+                height, width, tasks=None, task_names=None, data=None,
+                group_tasks=group_tasks
             )
             return fig
         else:
@@ -2367,14 +2455,14 @@ class FigureFactory(object):
                 fig = FigureFactory._gantt_colorscale(
                     chart, colors, title, index_col, show_colorbar, bar_width,
                     showgrid_x, showgrid_y, height, width,
-                    tasks=None, task_names=None, data=None
+                    tasks=None, task_names=None, data=None, group_tasks=group_tasks
                 )
                 return fig
             else:
                 fig = FigureFactory._gantt_dict(
                     chart, colors, title, index_col, show_colorbar, bar_width,
                     showgrid_x, showgrid_y, height, width,
-                    tasks=None, task_names=None, data=None
+                    tasks=None, task_names=None, data=None, group_tasks=group_tasks
                 )
                 return fig
 
@@ -2403,6 +2491,7 @@ class FigureFactory(object):
             else:
                 colors = list(colors)
 
+
         # convert color elements in list to tuple color
         for j, each_color in enumerate(colors):
             if 'rgb' in each_color:


Test patch:
diff --git a/plotly/tests/test_core/test_tools/test_figure_factory.py b/plotly/tests/test_core/test_tools/test_figure_factory.py
index 4fb488bfb..9ebc68c55 100644
--- a/plotly/tests/test_core/test_tools/test_figure_factory.py
+++ b/plotly/tests/test_core/test_tools/test_figure_factory.py
@@ -1269,6 +1269,181 @@ class TestGantt(TestCase):
                                 index_col='Complete',
                                 colors=['#ffffff'])
 
+
+    def test_gannt_groups_and_descriptions(self):
+
+        # check if grouped gantt chart matches with expected output
+
+        df = [
+            dict(Task='Task A', Description='Task A - 1', Start='2008-10-05',
+                 Finish='2009-04-15', IndexCol = 'TA'),
+            dict(Task="Task B", Description='Task B - 1', Start='2008-12-06',
+                 Finish='2009-03-15', IndexCol = 'TB'),
+            dict(Task="Task C", Description='Task C - 1', Start='2008-09-07',
+                 Finish='2009-03-15', IndexCol = 'TC'),
+            dict(Task="Task C", Description='Task C - 2', Start='2009-05-08',
+                 Finish='2009-04-15', IndexCol = 'TC'),
+            dict(Task="Task A", Description='Task A - 2', Start='2009-04-20',
+                 Finish='2009-05-30', IndexCol = 'TA')
+        ]
+
+        test_gantt_chart = tls.FigureFactory.create_gantt(
+            df, colors=dict(TA='rgb(220, 0, 0)', TB='rgb(170, 14, 200)',
+            TC=(1, 0.9, 0.16)), show_colorbar=True, index_col='IndexCol',
+            group_tasks=True
+        )
+
+        exp_gantt_chart = {
+            'data': [{'marker': {'color': 'white'},
+               'name': '',
+               'text': 'Task A - 1',
+               'x': ['2008-10-05', '2009-04-15'],
+               'y': [2, 2]},
+              {'marker': {'color': 'white'},
+               'name': '',
+               'text': 'Task B - 1',
+               'x': ['2008-12-06', '2009-03-15'],
+               'y': [1, 1]},
+              {'marker': {'color': 'white'},
+               'name': '',
+               'text': 'Task C - 1',
+               'x': ['2008-09-07', '2009-03-15'],
+               'y': [0, 0]},
+              {'marker': {'color': 'white'},
+               'name': '',
+               'text': 'Task C - 2',
+               'x': ['2009-05-08', '2009-04-15'],
+               'y': [0, 0]},
+              {'marker': {'color': 'white'},
+               'name': '',
+               'text': 'Task A - 2',
+               'x': ['2009-04-20', '2009-05-30'],
+               'y': [2, 2]},
+              {'hoverinfo': 'none',
+               'marker': {'color': 'rgb(220, 0, 0)', 'size': 1},
+               'name': 'TA',
+               'showlegend': True,
+               'x': ['2009-04-20', '2009-04-20'],
+               'y': [0, 0]},
+              {'hoverinfo': 'none',
+               'marker': {'color': 'rgb(170, 14, 200)', 'size': 1},
+               'name': 'TB',
+               'showlegend': True,
+               'x': ['2009-04-20', '2009-04-20'],
+               'y': [1, 1]},
+              {'hoverinfo': 'none',
+               'marker': {'color': 'rgb(255, 230, 41)', 'size': 1},
+               'name': 'TC',
+               'showlegend': True,
+               'x': ['2009-04-20', '2009-04-20'],
+               'y': [2, 2]}],
+            'layout': {'height': 600,
+                'hovermode': 'closest',
+                'shapes': [{'fillcolor': 'rgb(220, 0, 0)',
+                            'line': {'width': 0},
+                            'opacity': 1,
+                            'type': 'rect',
+                            'x0': '2008-10-05',
+                            'x1': '2009-04-15',
+                            'xref': 'x',
+                            'y0': 1.8,
+                            'y1': 2.2,
+                            'yref': 'y'},
+                           {'fillcolor': 'rgb(170, 14, 200)',
+                            'line': {'width': 0},
+                            'opacity': 1,
+                            'type': 'rect',
+                            'x0': '2008-12-06',
+                            'x1': '2009-03-15',
+                            'xref': 'x',
+                            'y0': 0.8,
+                            'y1': 1.2,
+                            'yref': 'y'},
+                           {'fillcolor': 'rgb(255, 230, 41)',
+                            'line': {'width': 0},
+                            'opacity': 1,
+                            'type': 'rect',
+                            'x0': '2008-09-07',
+                            'x1': '2009-03-15',
+                            'xref': 'x',
+                            'y0': -0.2,
+                            'y1': 0.2,
+                            'yref': 'y'},
+                           {'fillcolor': 'rgb(255, 230, 41)',
+                            'line': {'width': 0},
+                            'opacity': 1,
+                            'type': 'rect',
+                            'x0': '2009-05-08',
+                            'x1': '2009-04-15',
+                            'xref': 'x',
+                            'y0': -0.2,
+                            'y1': 0.2,
+                            'yref': 'y'},
+                           {'fillcolor': 'rgb(220, 0, 0)',
+                            'line': {'width': 0},
+                            'opacity': 1,
+                            'type': 'rect',
+                            'x0': '2009-04-20',
+                            'x1': '2009-05-30',
+                            'xref': 'x',
+                            'y0': 1.8,
+                            'y1': 2.2,
+                            'yref': 'y'}],
+                'showlegend': True,
+                'title': 'Gantt Chart',
+                'width': 900,
+                'xaxis': {'rangeselector': {'buttons': [{'count': 7,
+                                                         'label': '1w',
+                                                         'step': 'day',
+                                                         'stepmode': 'backward'},
+                                                        {'count': 1,
+                                                         'label': '1m',
+                                                         'step': 'month',
+                                                         'stepmode': 'backward'},
+                                                        {'count': 6,
+                                                         'label': '6m',
+                                                         'step': 'month',
+                                                         'stepmode': 'backward'},
+                                                        {'count': 1,
+                                                         'label': 'YTD',
+                                                         'step': 'year',
+                                                         'stepmode': 'todate'},
+                                                        {'count': 1,
+                                                         'label': '1y',
+                                                         'step': 'year',
+                                                         'stepmode': 'backward'},
+                                                        {'step': 'all'}]},
+                          'showgrid': False,
+                          'type': 'date',
+                          'zeroline': False},
+                'yaxis': {'autorange': False,
+                          'range': [-1, 4],
+                          'showgrid': False,
+                          'ticktext': ['Task C', 'Task B', 'Task A'],
+                          'tickvals': [0, 1, 2],
+                          'zeroline': False}}
+        }
+
+        self.assertEqual(test_gantt_chart['data'][0],
+                         exp_gantt_chart['data'][0])
+
+        self.assertEqual(test_gantt_chart['data'][1],
+                         exp_gantt_chart['data'][1])
+
+        self.assertEqual(test_gantt_chart['data'][2],
+                         exp_gantt_chart['data'][2])
+
+        self.assertEqual(test_gantt_chart['data'][3],
+                         exp_gantt_chart['data'][3])
+
+        self.assertEqual(test_gantt_chart['data'][4],
+                         exp_gantt_chart['data'][4])
+
+        self.assertEqual(test_gantt_chart['layout'],
+                         exp_gantt_chart['layout'])
+
+
+
     def test_gantt_all_args(self):
 
         # check if gantt chart matches with expected output