update implementation to use the astor library

.travis.yml

@@ -15,6 +15,6 @@ before_install:
 install:
   - python setup.py develop
 script:
-  - coverage run -m unittest2 discover
+  - coverage run -m haas
 after_success:
   - codecov

setup.py

@@ -110,6 +110,7 @@ def write_version_py(filename='trait_documenter/_version.py'):
         description='Autodoc extention for documenting traits',
         long_description=open('README.rst').read(),
         license="BSD",
+        install_requires=['sphinx', 'astor', 'traits'],
         classifiers=[
             "Programming Language :: Python :: 2.6",
             "Programming Language :: Python :: 2.7",

test-requirements.txt

@@ -1,3 +1,4 @@
 coverage
 mock
-unittest2
+unittest2 ; python_version < '2.7'
+haas

trait_documenter/tests/test_class_trait_documenter.py

@@ -82,7 +82,7 @@ def test_add_directive_header(self):
             ('.. py:attribute:: Dummy.trait_2', '<autodoc>'),
             ('   :noindex:', '<autodoc>'),
             ('   :module: trait_documenter.tests.test_file', '<autodoc>'),
-            ("   :annotation: = Property(Float,depends_on='trait_1')", '<autodoc>')]  # noqa
+            ("   :annotation: = Property(Float, depends_on='trait_1')", '<autodoc>')]  # noqa
         calls = documenter.add_line.call_args_list
         for index, line in enumerate(expected):
             self.assertEqual(calls[index][0], line)

trait_documenter/tests/test_file.py

@@ -6,9 +6,11 @@
     low=0.2,
     high=34)
 
+
 def dummy_function():
     pass
 
+
 class Dummy(HasTraits):
 
     trait_1 = Float
@@ -19,6 +21,12 @@ class Dummy(HasTraits):
 
     not_trait = 2
 
+    trait_3 = Property(
+        Float,  # first comment
+        depends_on='trait_4')
+
+    trait_4 = Float  # second comment
+
 
 class Dummy1(HasTraits):
 

trait_documenter/tests/test_get_trait_definition.py

@@ -2,6 +2,7 @@
 import unittest
 
 from trait_documenter.tests import test_file, test_file2
+from trait_documenter.tests.testing import expected_failure_when, is_python_26
 from trait_documenter.tests.test_file import Dummy, dummy_function
 from trait_documenter.util import get_trait_definition, DefinitionError
 
@@ -30,6 +31,7 @@ def test_get_inavlid_module_trait_definition(self):
             get_trait_definition,
             parent, object_name)
 
+    @expected_failure_when(is_python_26())
     def test_get_multi_line_module_trait_definition(self):
         # given
         parent = test_file
@@ -39,7 +41,7 @@ def test_get_multi_line_module_trait_definition(self):
         definition = get_trait_definition(parent, object_name)
 
         # then
-        self.assertEqual(definition, "Range(low=0.2,high=34)")
+        self.assertEqual(definition, "Range(low=0.2, high=34)")
 
     def test_get_simple_class_trait_definition(self):
         # given
@@ -62,6 +64,17 @@ def test_get_trait_definition_from_function(self):
             get_trait_definition,
             dummy_function, object_name)
 
+    def test_get_simple_class_trait_definition_with_comments(self):
+        # given
+        parent = Dummy
+        object_name = 'trait_4'
+
+        # when
+        definition = get_trait_definition(parent, object_name)
+
+        # then
+        self.assertEqual(definition, 'Float')
+
     def test_get_multi_line_class_trait_definition(self):
         # given
         parent = Dummy
@@ -71,7 +84,7 @@ def test_get_multi_line_class_trait_definition(self):
         definition = get_trait_definition(parent, object_name)
 
         # then
-        self.assertEqual(definition, "Property(Float,depends_on='trait_1')")
+        self.assertEqual(definition, "Property(Float, depends_on='trait_1')")
 
     def test_get_trait_definition_over_horrible_assigment(self):
         # given
@@ -95,6 +108,17 @@ def test_get_trait_definition_inside_if_block(self):
         # then
         self.assertEqual(definition, "List(Int)")
 
+    def test_get_multi_line_class_trait_definition_with_comment(self):
+        # given
+        parent = Dummy
+        object_name = 'trait_3'
+
+        # when
+        definition = get_trait_definition(parent, object_name)
+
+        # then
+        self.assertEqual(definition, "Property(Float, depends_on='trait_4')")
+
 
 if __name__ == '__main__':
     unittest.main()

trait_documenter/tests/test_module_trait_documenter.old

@@ -0,0 +1,46 @@
+from __future__ import unicode_literals
+import unittest
+
+import mock
+
+from trait_documenter.module_trait_documenter import (
+    ModuleTraitDocumenter)
+from trait_documenter.tests import test_file
+from trait_documenter.tests.test_file import Dummy
+
+
+class TestModuleTraitDocumenter(unittest.TestCase):
+
+    def test_can_document_member(self):
+        can_document_member = ModuleTraitDocumenter.can_document_member
+        parent = mock.Mock()
+
+        # modules
+        parent.object = test_file
+        self.assertFalse(can_document_member(Dummy, 'Dummy', False, parent))
+        self.assertFalse(can_document_member(
+            test_file.module_trait, 'module_trait', False, parent))
+        self.assertTrue(can_document_member(
+            test_file.module_trait, 'module_trait', True, parent))
+
+        parent.object = Dummy
+        self.assertFalse(
+            can_document_member(
+                Dummy.class_traits()['trait_1'], 'trait_1', True, parent))
+        self.assertFalse(
+            can_document_member(
+                Dummy.class_traits()['trait_1'], 'trait_1', False, parent))
+        self.assertFalse(
+            can_document_member(
+                Dummy.not_trait, 'not_trait', False, parent))
+
+    def test_module_level_trait_definition(self):
+        documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
+        documenter.parent = test_file
+        documenter.object_name = 'module_trait'
+        definition = documenter.get_trait_definition()
+        self.assertEqual(definition, "Int")
+
+
+if __name__ == '__main__':
+    unittest.main()

trait_documenter/tests/test_module_trait_documenter.py

@@ -6,6 +6,7 @@
 
 from trait_documenter.module_trait_documenter import ModuleTraitDocumenter
 from trait_documenter.tests import test_file
+from trait_documenter.tests.testing import expected_failure_when, is_python_26
 from trait_documenter.tests.test_file import Dummy
 
 
@@ -47,6 +48,7 @@ def test_import_object(self):
         self.assertTrue(documenter.object is not None)
         self.assertEqual(documenter.parent, test_file)
 
+    @expected_failure_when(is_python_26())
     def test_add_directive_header(self):
         # given
         documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
@@ -66,7 +68,7 @@ def test_add_directive_header(self):
             ('.. py:data:: long_module_trait', '<autodoc>'),
             ('   :noindex:', '<autodoc>'),
             ('   :module: trait_documenter.tests.test_file', '<autodoc>'),
-            ('   :annotation: = Range(low=0.2,high=34)', '<autodoc>')]
+            ('   :annotation: = Range(low=0.2, high=34)', '<autodoc>')]
         calls = documenter.add_line.call_args_list
         for index, line in enumerate(expected):
             self.assertEqual(calls[index][0], line)

trait_documenter/tests/testing.py

@@ -0,0 +1,19 @@
+def is_python_26():
+    import sys
+    return sys.version_info < (2, 6)
+
+
+class expected_failure_when(object):
+
+    def __init__(self, condition):
+        self.condition = condition
+
+    def __call__(self, function):
+        if self.condition:
+            if is_python_26():
+                import unittest2 as unittest
+            else:
+                import unittest
+            return unittest.expectedFailure(function)
+        else:
+            return function

trait_documenter/util.py

@@ -3,6 +3,8 @@
 import collections
 from _ast import ClassDef, Assign, Name
 
+import astor
+
 
 class DefinitionError(Exception):
     pass
@@ -57,17 +59,7 @@ def get_trait_definition(parent, trait_name):
         # we always get the last assignment in the file
         node, name = assignments[-1]
 
-    endlineno = name.lineno
-    for item in ast.walk(node):
-        if hasattr(item, 'lineno'):
-            endlineno = max(endlineno, item.lineno)
-
-    definition_lines = [
-        line.strip()
-        for line in source.splitlines()[name.lineno-1:endlineno]]
-    definition = ''.join(definition_lines)
-    equal = definition.index('=')
-    return definition[equal + 1:].lstrip()
+    return astor.to_source(node.value).strip()
 
 
 def trait_node(node, trait_name):