Fix trait_documenter to properly document traits of base classes

traits/util/tests/test_trait_documenter.py

@@ -100,21 +100,27 @@ def not_a_trait(self):
         """
 
 
-class MySubClass(MyTestClass):
+class MySubClassAppend(MyTestClass):
 
     #: A new attribute.
     foo = Bool(True)
 
 
+class MySubClassReplace(MyTestClass):
+
+    #: Replace attribute.
+    bar = Int(1)
+
+
 @requires_sphinx
 class TestTraitDocumenter(unittest.TestCase):
     """ Tests for the trait documenter. """
 
     def setUp(self):
         self.source = """
     depth_interval = Property(Tuple(Float, Float),
-                              depends_on="_depth_interval")
-"""
+                              depends_on="_depth_interval")\n
+    """
         string_io = io.StringIO(self.source)
         tokens = tokenize.generate_tokens(string_io.readline)
         self.tokens = tokens
@@ -136,24 +142,6 @@ def test_get_definition_tokens(self):
 
         self.assertEqual(src.rstrip(), string)
 
-    def test_add_line(self):
-        mocked_directive = mock.MagicMock()
-
-        documenter = TraitDocumenter(mocked_directive, "test", "   ")
-        documenter.object_name = "test_attribute"
-        documenter.parent = Fake
-
-        with mock.patch(
-            (
-                "traits.util.trait_documenter.ClassLevelDocumenter"
-                ".add_directive_header"
-            )
-        ):
-            documenter.add_directive_header("")
-
-        self.assertEqual(
-            len(documenter.directive.result.append.mock_calls), 1)
-
     def test_abbreviated_annotations(self):
         # Regression test for enthought/traits#493.
         with self.create_directive() as directive:
@@ -243,12 +231,12 @@ def test_class(self):
         for index, line in enumerate(expected):
             self.assertEqual(calls[index][0], line)
 
-    def test_subclass(self):
+    def test_subclass_append(self):
         # given
         documenter = TraitDocumenter(mock.Mock(), 'test')
         documenter.object_name = 'bar'
-        documenter.objpath = ['MySubClass', 'bar']
-        documenter.parent = MySubClass
+        documenter.objpath = ['MySubClassAppend', 'bar']
+        documenter.parent = MySubClassAppend
         documenter.modname = 'traits.util.tests.test_trait_documenter'
         documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
         documenter.add_line = mock.Mock()
@@ -259,7 +247,7 @@ def test_subclass(self):
         # then
         self.assertEqual(documenter.directive.warn.call_args_list, [])
         expected = [
-            ('.. py:attribute:: MySubClass.bar', '<autodoc>'),
+            ('.. py:attribute:: MySubClassAppend.bar', '<autodoc>'),
             (f'   :{no_index}:', '<autodoc>'),
             ('   :module: traits.util.tests.test_trait_documenter', '<autodoc>'),  # noqa
             ('   :annotation: = Int(42, desc=""" First line …', '<autodoc>')]  # noqa
@@ -271,7 +259,33 @@ def test_subclass(self):
 
         # given
         documenter.object_name = 'foo'
-        documenter.objpath = ['MySubClass', 'foo']
+        documenter.objpath = ['MySubClassAppend', 'foo']
+        documenter.add_line = mock.Mock()
+
+        # when
+        documenter.add_directive_header('')
+
+        # then
+        self.assertEqual(documenter.directive.warn.call_args_list, [])
+        expected = [
+            ('.. py:attribute:: MySubClassAppend.foo', '<autodoc>'),
+            (f'   :{no_index}:', '<autodoc>'),
+            ('   :module: traits.util.tests.test_trait_documenter', '<autodoc>'),  # noqa
+            ('   :annotation: = Bool(True)', '<autodoc>')]
+        if no_index_entry:
+            expected.insert(2, ('   :no-index-entry:', '<autodoc>'))
+        calls = documenter.add_line.call_args_list
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
+
+    def test_subclass_replace(self):
+        # given
+        documenter = TraitDocumenter(mock.Mock(), 'test')
+        documenter.object_name = 'bar'
+        documenter.objpath = ['MySubClassReplace', 'bar']
+        documenter.parent = MySubClassReplace
+        documenter.modname = 'traits.util.tests.test_trait_documenter'
+        documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
         documenter.add_line = mock.Mock()
 
         # when
@@ -280,10 +294,10 @@ def test_subclass(self):
         # then
         self.assertEqual(documenter.directive.warn.call_args_list, [])
         expected = [
-            ('.. py:attribute:: MySubClass.foo', '<autodoc>'),
+            ('.. py:attribute:: MySubClassReplace.bar', '<autodoc>'),
             (f'   :{no_index}:', '<autodoc>'),
             ('   :module: traits.util.tests.test_trait_documenter', '<autodoc>'),  # noqa
-            ('   :annotation: = Bool(True)', '<autodoc>')]  # noqa
+            ('   :annotation: = Int(1)', '<autodoc>')]  # noqa
         if no_index_entry:
             expected.insert(2, ('   :no-index-entry:', '<autodoc>'))
         calls = documenter.add_line.call_args_list

traits/util/trait_documenter.py

@@ -10,18 +10,17 @@
 
 """
     A Trait Documenter
-    (Subclassed from the autodoc ClassLevelDocumenter)
+    (Subclassed from the autodoc AttributeDocumenter)
 
 """
-from importlib import import_module
 import inspect
 import io
 import types
 import token
 import tokenize
 import traceback
 
-from sphinx.ext.autodoc import ClassLevelDocumenter
+from sphinx.ext.autodoc import AttributeDocumenter, import_module
 from sphinx.util import logging
 
 from traits.has_traits import MetaHasTraits
@@ -42,7 +41,7 @@ def _is_class_trait(name, cls):
     )
 
 
-class TraitDocumenter(ClassLevelDocumenter):
+class TraitDocumenter(AttributeDocumenter):
     """ Specialized Documenter subclass for trait attributes.
 
     The class defines a new documenter that recovers the trait definition
@@ -114,7 +113,6 @@ def add_directive_header(self, sig):
         option set to the trait definition.
 
         """
-        ClassLevelDocumenter.add_directive_header(self, sig)
         # Look into the class and parent classes:
         parent = self.parent
         classes = list(types.resolve_bases(parent.__bases__))
@@ -138,8 +136,8 @@ def add_directive_header(self, sig):
         # throw away all lines after the first.
         if "\n" in definition:
             definition = definition.partition("\n")[0] + " …"
-
-        self.add_line("   :annotation: = {0}".format(definition), "<autodoc>")
+        self.options.annotation = f'= {definition}'
+        super().add_directive_header(sig)
 
 
 def trait_definition(*, cls, trait_name):
@@ -179,23 +177,15 @@ class MyModel(HasStrictTraits)
     tokens = tokenize.generate_tokens(string_io.readline)
 
     # find the trait definition start
-    trait_found = False
     name_found = False
-    while not trait_found:
-        item = next(tokens, None)
-        if item is None:
-            break
+    for item in tokens:
         if name_found and item[:2] == (token.OP, "="):
-            trait_found = True
-            continue
+            break
         if item[:2] == (token.NAME, trait_name):
             name_found = True
-
-    if not trait_found:
-        raise ValueError(
-            "No trait definition for {!r} found in {!r}".format(
-                trait_name, cls)
-        )
+    else:
+        message = "No trait definition for {!r} found in {!r}"
+        raise ValueError(message.format(trait_name, cls))
 
     # Retrieve the trait definition.
     definition_tokens = _get_definition_tokens(tokens)
@@ -208,8 +198,8 @@ def _get_definition_tokens(tokens):
 
     Parameters
     ----------
-    tokens : iterator
-        An iterator producing tokens.
+    tokens : iteratable
+        An iteratable producing tokens.
 
     Returns
     -------
@@ -235,10 +225,10 @@ def _get_definition_tokens(tokens):
         )
 
         definition_tokens.append(item)
-
     return definition_tokens
 
 
 def setup(app):
     """ Add the TraitDocumenter in the current sphinx autodoc instance. """
+    app.setup_extension('sphinx.ext.autodoc')  # Require autodoc extension
     app.add_autodocumenter(TraitDocumenter)