From d86ea47b52b1df6d578cad88efd509802234f37a Mon Sep 17 00:00:00 2001
From: Takeshi KOMIYA <i.tkomiya@gmail.com>
Date: Fri, 11 Mar 2016 13:20:05 +0900
Subject: [PATCH] Add ``:caption:`` option for sphinx.ext.inheritance_diagram

---
 CHANGES                                       |  2 ++
 doc/ext/inheritance.rst                       |  5 +++++
 sphinx/ext/inheritance_diagram.py             |  8 ++++++-
 .../test-ext-inheritance_diagram/index.rst    |  3 +++
 tests/test_ext_inheritance_diagram.py         | 22 +++++++++++++++++++
 5 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/CHANGES b/CHANGES
index 20beb2475..52bbc6682 100644
--- a/CHANGES
+++ b/CHANGES
@@ -7,6 +7,8 @@ Incompatible changes
 Features added
 --------------
 
+* Add ``:caption:`` option for sphinx.ext.inheritance_diagram.
+
 Bugs fixed
 ----------
 
diff --git a/doc/ext/inheritance.rst b/doc/ext/inheritance.rst
index 5e0a76fcc..dd8d5aa99 100644
--- a/doc/ext/inheritance.rst
+++ b/doc/ext/inheritance.rst
@@ -33,10 +33,15 @@ It adds this directive:
    It also supports a ``private-bases`` flag option; if given, private base
    classes (those whose name starts with ``_``) will be included.
 
+   You can use ``caption`` option to give a caption to the diagram.
+
    .. versionchanged:: 1.1
       Added ``private-bases`` option; previously, all bases were always
       included.
 
+   .. versionchanged:: 1.5
+      Added ``caption`` option 
+
 
 New config values are:
 
diff --git a/sphinx/ext/inheritance_diagram.py b/sphinx/ext/inheritance_diagram.py
index a06c4b17c..11af67dc5 100644
--- a/sphinx/ext/inheritance_diagram.py
+++ b/sphinx/ext/inheritance_diagram.py
@@ -52,7 +52,7 @@ from docutils.parsers.rst import directives
 
 import sphinx
 from sphinx.ext.graphviz import render_dot_html, render_dot_latex, \
-    render_dot_texinfo
+    render_dot_texinfo, figure_wrapper
 from sphinx.pycode import ModuleAnalyzer
 from sphinx.util import force_decode
 from sphinx.util.compat import Directive
@@ -297,6 +297,7 @@ class InheritanceDiagram(Directive):
     option_spec = {
         'parts': directives.nonnegative_int,
         'private-bases': directives.flag,
+        'caption': directives.unchanged,
     }
 
     def run(self):
@@ -330,6 +331,11 @@ class InheritanceDiagram(Directive):
         # Store the graph object so we can use it to generate the
         # dot file later
         node['graph'] = graph
+
+        # wrap the result in figure node
+        caption = self.options.get('caption')
+        if caption:
+            node = figure_wrapper(self, node, caption)
         return [node]
 
 
diff --git a/tests/roots/test-ext-inheritance_diagram/index.rst b/tests/roots/test-ext-inheritance_diagram/index.rst
index 876996ca8..777192bd7 100644
--- a/tests/roots/test-ext-inheritance_diagram/index.rst
+++ b/tests/roots/test-ext-inheritance_diagram/index.rst
@@ -3,3 +3,6 @@ test-ext-inheritance_diagram
 ============================
 
 .. inheritance-diagram:: test.Foo
+
+.. inheritance-diagram:: test.Foo
+   :caption: Test Foo!
diff --git a/tests/test_ext_inheritance_diagram.py b/tests/test_ext_inheritance_diagram.py
index 64446eed8..89925f754 100644
--- a/tests/test_ext_inheritance_diagram.py
+++ b/tests/test_ext_inheritance_diagram.py
@@ -9,9 +9,31 @@
     :license: BSD, see LICENSE for details.
 """
 
+import re
 from util import with_app
 
 
 @with_app('html', testroot='ext-inheritance_diagram')
 def test_inheritance_diagram_html(app, status, warning):
     app.builder.build_all()
+
+    content = (app.outdir / 'index.html').text()
+
+    pattern = ('<div class="figure" id="id1">\n'
+               '<img src="_images/inheritance-\w+.png" alt="Inheritance diagram of test.Foo" '
+               'class="inheritance"/>\n<p class="caption"><span class="caption-text">'
+               'Test Foo!</span><a class="headerlink" href="#id1" '
+               'title="Permalink to this image">\xb6</a></p>')
+    assert re.search(pattern, content, re.M)
+
+
+@with_app('latex', testroot='ext-inheritance_diagram')
+def test_inheritance_diagram_latex(app, status, warning):
+    app.builder.build_all()
+
+    content = (app.outdir / 'Python.tex').text()
+
+    pattern = ('\\\\begin{figure}\\[htbp]\n\\\\centering\n\\\\capstart\n\n'
+               '\\\\includegraphics{inheritance-\\w+.pdf}\n'
+               '\\\\caption{Test Foo!}\\\\label{index:id1}\\\\end{figure}')
+    assert re.search(pattern, content, re.M)