Merge pull request #9081 from fariza/table-array-like

QuLogic · web-flow · commit 5a75f5fd74b1 · 2018-02-07T17:58:07.000-05:00
cell returned when added to Table
diff --git a/doc/api/index.rst b/doc/api/index.rst
@@ -49,6 +49,7 @@
    scale_api.rst
    spines_api.rst
    style_api.rst
+   table_api.rst
    text_api.rst
    ticker_api.rst
    tight_layout_api.rst
diff --git a/doc/api/table_api.rst b/doc/api/table_api.rst
@@ -0,0 +1,11 @@
+*****
+table
+*****
+
+:mod:`matplotlib.table`
+===================================
+
+.. automodule:: matplotlib.table
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -2021,10 +2021,17 @@ def _update_patch_limits(self, patch):
 
     def add_table(self, tab):
         """
-        Add a :class:`~matplotlib.tables.Table` instance to the
+        Add a :class:`~matplotlib.table.Table` instance to the
         list of axes tables
 
-        Returns the table.
+        Parameters
+        ----------
+        tab: `matplotlib.table.Table`
+            Table instance
+
+        Returns
+        -------
+        `matplotlib.table.Table`: the table.
         """
         self._set_artist_props(tab)
         self.tables.append(tab)
diff --git a/lib/matplotlib/table.py b/lib/matplotlib/table.py
@@ -276,17 +276,52 @@ def __init__(self, ax, loc=None, bbox=None, **kwargs):
         self.set_clip_on(False)
 
     def add_cell(self, row, col, *args, **kwargs):
-        """ Add a cell to the table. """
-        xy = (0, 0)
+        """
+        Add a cell to the table.
+
+        Parameters
+        ----------
+        row : int
+            Row index
+        col : int
+            Column index
+
+        Returns
+        -------
+        `CustomCell`: Automatically created cell
 
+        """
+        xy = (0, 0)
         cell = CustomCell(xy, visible_edges=self.edges, *args, **kwargs)
+        self[row, col] = cell
+        return cell
+
+    def __setitem__(self, position, cell):
+        """
+        Set a customcell in a given position
+        """
+        if not isinstance(cell, CustomCell):
+            raise TypeError('Table only accepts CustomCell')
+        try:
+            row, col = position[0], position[1]
+        except Exception:
+            raise KeyError('Only tuples length 2 are accepted as coordinates')
         cell.set_figure(self.figure)
         cell.set_transform(self.get_transform())
-
         cell.set_clip_on(False)
         self._cells[row, col] = cell
         self.stale = True
 
+    def __getitem__(self, position):
+        """
+        Retreive a custom cell from a given position
+        """
+        try:
+            row, col = position[0], position[1]
+        except Exception:
+            raise KeyError('Only tuples length 2 are accepted as coordinates')
+        return self._cells[row, col]
+
     @property
     def edges(self):
         return self._edges
diff --git a/lib/matplotlib/tests/test_table.py b/lib/matplotlib/tests/test_table.py
@@ -6,7 +6,7 @@
 import numpy as np
 from matplotlib.testing.decorators import image_comparison
 
-from matplotlib.table import CustomCell
+from matplotlib.table import CustomCell, Table
 from matplotlib.path import Path
 
 
@@ -181,3 +181,16 @@ def test_auto_column():
     tb4.auto_set_font_size(False)
     tb4.set_fontsize(12)
     tb4.auto_set_column_width("-101")
+
+
+def test_table_cells():
+    fig, ax = plt.subplots()
+    table = Table(ax)
+
+    cell = table.add_cell(1, 2, 1, 1)
+    assert isinstance(cell, CustomCell)
+    assert cell is table[1, 2]
+
+    cell2 = CustomCell((0, 0), 1, 2, visible_edges=None)
+    table[2, 1] = cell2
+    assert table[2, 1] is cell2
