fix annotation formatting for builtin types in Python 2.x

2 files changed, 28 insertions, 1 deletions

diff --git a/funcsigs/__init__.py b/funcsigs/__init__.py
index f2924c3..b9ba326 100644
--- a/funcsigs/__init__.py
+++ b/funcsigs/__init__.py
@@ -30,7 +30,7 @@ _NonUserDefinedCallables = (_WrapperDescriptor,
 
 def formatannotation(annotation, base_module=None):
     if isinstance(annotation, type):
-        if annotation.__module__ in ('builtins', base_module):
+        if annotation.__module__ in ('builtins', '__builtin__', base_module):
             return annotation.__name__
         return annotation.__module__+'.'+annotation.__name__
     return repr(annotation)
diff --git a/tests/test_formatannotation.py b/tests/test_formatannotation.py
new file mode 100644
index 0000000..fd7a887
--- /dev/null
+++ b/tests/test_formatannotation.py
@@ -0,0 +1,27 @@
+try:
+    # python 2.x
+    import unittest2 as unittest
+except ImportError:
+    # python 3.x
+    import unittest
+
+import funcsigs
+
+
+class TestFormatAnnotation(unittest.TestCase):
+    def test_string (self):
+        self.assertEqual(funcsigs.formatannotation("annotation"),
+                         "'annotation'")
+
+    def test_builtin_type (self):
+        self.assertEqual(funcsigs.formatannotation(int),
+                         "int")
+
+    def test_user_type (self):
+        class dummy (object): pass
+        self.assertEqual(funcsigs.formatannotation(dummy),
+                         "tests.test_formatannotation.dummy")
+
+
+if __name__ == "__main__":
+    unittest.begin()