Fix trait_documenter to be less fragile

traits/util/tests/test_trait_documenter.py

@@ -20,7 +20,7 @@
 import unittest
 import unittest.mock as mock
 
-from traits.api import HasTraits, Int
+from traits.api import Bool, HasTraits, Int, Property
 from traits.testing.optional_dependencies import sphinx, requires_sphinx
 
 
@@ -33,6 +33,7 @@
 
     from traits.util.trait_documenter import (
         _get_definition_tokens,
+        trait_definition,
         TraitDocumenter,
     )
 
@@ -65,6 +66,12 @@ class MyTestClass(HasTraits):
     """)
 
 
+class Fake(HasTraits):
+
+    #: Test attribute
+    test_attribute = Property(Bool, label="ミスあり")
+
+
 class FindTheTraits(HasTraits):
     """
     Class for testing the can_document_member functionality.
@@ -118,29 +125,19 @@ def test_get_definition_tokens(self):
 
     def test_add_line(self):
 
-        src = textwrap.dedent(
-            """\
-        class Fake(HasTraits):
-
-            #: Test attribute
-            test = Property(Bool, label="ミスあり")
-        """
-        )
         mocked_directive = mock.MagicMock()
 
         documenter = TraitDocumenter(mocked_directive, "test", "   ")
-        documenter.object_name = "Property"
+        documenter.object_name = "test_attribute"
+        documenter.parent = Fake
 
         with mock.patch(
-            "traits.util.trait_documenter.inspect.getsource", return_value=src
+            (
+                "traits.util.trait_documenter.ClassLevelDocumenter"
+                ".add_directive_header"
+            )
         ):
-            with mock.patch(
-                (
-                    "traits.util.trait_documenter.ClassLevelDocumenter"
-                    ".add_directive_header"
-                )
-            ):
-                documenter.add_directive_header("")
+            documenter.add_directive_header("")
 
         self.assertEqual(
             len(documenter.directive.result.append.mock_calls), 1)
@@ -164,6 +161,16 @@ def test_abbreviated_annotations(self):
         self.assertIn("First line", item)
         self.assertNotIn("\n", item)
 
+    def test_successful_trait_definition(self):
+        definition = trait_definition(cls=Fake, trait_name="test_attribute")
+        self.assertEqual(
+            definition, 'Property(Bool, label="ミスあり")',
+        )
+
+    def test_failed_trait_definition(self):
+        with self.assertRaises(ValueError):
+            trait_definition(cls=Fake, trait_name="not_a_trait")
+
     def test_can_document_member(self):
         # Regression test for enthought/traits#1238
 

traits/util/trait_documenter.py

@@ -21,12 +21,16 @@
 import traceback
 
 from sphinx.ext.autodoc import ClassLevelDocumenter
+from sphinx.util import logging
 
 from traits.has_traits import MetaHasTraits
 from traits.trait_type import TraitType
 from traits.traits import generic_trait
 
 
+logger = logging.getLogger(__name__)
+
+
 def _is_class_trait(name, cls):
     """ Check if the name is in the list of class defined traits of ``cls``.
     """
@@ -120,7 +124,19 @@ def add_directive_header(self, sig):
 
         """
         ClassLevelDocumenter.add_directive_header(self, sig)
-        definition = self._get_trait_definition()
+        try:
+            definition = trait_definition(
+                cls=self.parent,
+                trait_name=self.object_name,
+            )
+        except ValueError:
+            # Without this, a failure to find the trait definition aborts
+            # the whole documentation build.
+            logger.warn(
+                "No definition for the trait {!r} could be found in "
+                "class {!r}.".format(self.object_name, self.parent),
+                exc_info=True)
+            return
 
         # Workaround for enthought/traits#493: if the definition is multiline,
         # throw away all lines after the first.
@@ -129,32 +145,66 @@ def add_directive_header(self, sig):
 
         self.add_line("   :annotation: = {0}".format(definition), "<autodoc>")
 
-    # Private Interface #####################################################
 
-    def _get_trait_definition(self):
-        """ Retrieve the Trait attribute definition
-        """
+def trait_definition(*, cls, trait_name):
+    """ Retrieve the portion of the source defining a Trait attribute.
+
+    For example, given a class::
+
+        class MyModel(HasStrictTraits)
+            foo = List(Int, [1, 2, 3])
+
+    ``trait_definition(cls=MyModel, trait_name="foo")`` returns
+    ``"List(Int, [1, 2, 3])"``.
+
+    Parameters
+    ----------
+    cls : MetaHasTraits
+        Class being documented.
+    trait_name : str
+        Name of the trait being documented.
 
-        # Get the class source and tokenize it.
-        source = inspect.getsource(self.parent)
-        string_io = io.StringIO(source)
-        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)
-            if name_found and item[:2] == (token.OP, "="):
-                trait_found = True
-                continue
-            if item[:2] == (token.NAME, self.object_name):
-                name_found = True
-
-        # Retrieve the trait definition.
-        definition_tokens = _get_definition_tokens(tokens)
-        definition = tokenize.untokenize(definition_tokens).strip()
-        return definition
+    Returns
+    -------
+    str
+        The portion of the source containing the trait definition. For
+        example, for a class trait defined as ``"my_trait = Float(3.5)"``,
+        the returned string will contain ``"Float(3.5)"``.
+
+    Raises
+    ------
+    ValueError
+        If *trait_name* doesn't appear as a class-level variable in the
+        source.
+    """
+    # Get the class source and tokenize it.
+    source = inspect.getsource(cls)
+    string_io = io.StringIO(source)
+    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
+        if name_found and item[:2] == (token.OP, "="):
+            trait_found = True
+            continue
+        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)
+        )
+
+    # Retrieve the trait definition.
+    definition_tokens = _get_definition_tokens(tokens)
+    definition = tokenize.untokenize(definition_tokens).strip()
+    return definition
 
 
 def _get_definition_tokens(tokens):