Support handling nested classes (#85)

Without understanding the module hierarchy, we can't distinguish
between a module and a class in a dotted name like `foo.bar.A.B`, and
will insert code like `from foo.bar.A import B`, even though `A` is a
class.

Work around this by using : to separate the module and class name
if the class is nested inside another class.

The fix works only on Python 3 since it uses `__qualname__`.

Author: msullivan

pyannotate_runtime/collect_types.py

@@ -54,7 +54,7 @@
 )
 from contextlib import contextmanager
 
-MYPY=False  
+MYPY=False
 if MYPY:
     # MYPY is True when mypy is running
     # 'Type' is only required for running mypy, not for running pyannotate
@@ -430,7 +430,9 @@ def name_from_type(type_):
                 # Also ignore '<uknown>' modules so pyannotate can parse these types
                 return type_.__name__
             else:
-                return '%s.%s' % (module, type_.__name__)
+                name = getattr(type_, '__qualname__', None) or type_.__name__
+                delim = '.' if '.' not in name else ':'
+                return '%s%s%s' % (module, delim, name)
         else:
             return 'None'
 

pyannotate_runtime/tests/test_collect_types.py

@@ -9,6 +9,7 @@
 import json
 import os
 import sched
+import sys
 import time
 import unittest
 from collections import namedtuple
@@ -51,14 +52,18 @@ def noop_dec(a):
     # type: (Any) -> Any
     return a
 
+def discard(a):
+    # type: (Any) -> None
+    pass
 
 @noop_dec
 class FoosParent(object):
     pass
 
 
 class FooObject(FoosParent):
-    pass
+    class FooNested(object):
+        pass
 
 
 class FooReturn(FoosParent):
@@ -314,6 +319,8 @@ def test_run_a_bunch_of_tests(self):
 
             OldStyleClass().foo(10)
 
+            discard(FooObject.FooNested())
+
         # TODO(svorobev): add checks for the rest of the functions
         # print_int,
         self.assert_type_comments(
@@ -324,6 +331,12 @@ def test_run_a_bunch_of_tests(self):
             ['(int, pyannotate_runtime.tests.test_collect_types.FooNamedTuple) -> EOFError'])
         self.assert_type_comments('OldStyleClass.foo', ['(int) -> int'])
 
+        # Need __qualname__ to get this right
+        if sys.version_info >= (3, 3):
+            self.assert_type_comments(
+                'discard',
+                ['(pyannotate_runtime.tests.test_collect_types:FooObject.FooNested) -> None'])
+
         # TODO: that could be better
         self.assert_type_comments('takes_different_lists', ['(List[Union[int, str]]) -> None'])
 

pyannotate_tools/annotations/parse.py

@@ -188,7 +188,7 @@ def tokenize(s):
             tokens.append(Separator('->'))
             s = s[2:]
         else:
-            m = re.match(r'[-\w]+( *\. *[-/\w]*)*', s)
+            m = re.match(r'[-\w]+(\s*(\.|:)\s*[-/\w]*)*', s)
             if not m:
                 raise ParseError(original)
             fullname = m.group(0)

pyannotate_tools/fixes/fix_annotate_json.py

@@ -266,19 +266,27 @@ def get_annotation_from_stub(self, node, results, funcname):
     def update_type_names(self, type_str):
         # Replace e.g. `List[pkg.mod.SomeClass]` with
         # `List[SomeClass]` and remember to import it.
-        return re.sub(r'[\w.]+', self.type_updater, type_str)
+        return re.sub(r'[\w.:]+', self.type_updater, type_str)
 
     def type_updater(self, match):
         # Replace `pkg.mod.SomeClass` with `SomeClass`
         # and remember to import it.
         word = match.group()
         if word == '...':
             return word
-        if '.' not in word:
+        if '.' not in word and ':' not in word:
             # Assume it's either builtin or from `typing`
             if word in typing_all:
                 self.add_import('typing', word)
             return word
-        mod, name = word.rsplit('.', 1)
-        self.add_import(mod, name)
+        # If there is a :, treat that as the separator between the
+        # module and the class.  Otherwise assume everything but the
+        # last element is the module.
+        if ':' in word:
+            mod, name = word.split(':')
+            to_import = name.split('.', 1)[0]
+        else:
+            mod, name = word.rsplit('.', 1)
+            to_import = name
+        self.add_import(mod, to_import)
         return name

pyannotate_tools/fixes/tests/test_annotate_json_py2.py

@@ -660,3 +660,24 @@ def nop(a):
                 return 0
             """
         self.check(a, b)
+
+    def test_nested(self):
+        self.setTestData(
+            [{"func_name": "nop",
+              "path": "<string>",
+              "line": 1,
+              "signature": {
+                  "arg_types": ["foo:A.B"],
+                  "return_type": "None"},
+              }])
+        a = """\
+            def nop(a):
+                pass
+            """
+        b = """\
+            from foo import A
+            def nop(a):
+                # type: (A.B) -> None
+                pass
+            """
+        self.check(a, b)