Creates "from __main__" import for __main__ program

[blueyed]

Given t-pyannotate.py:

from pyannotate_runtime import collect_types

collect_types.init_types_collection()
collect_types.start()

class G:
    pass

def c(obj):
    pass

c(G())

collect_types.stop()
collect_types.dump_stats("type_info.json")

Running and writing it produces the following diff, where it tries to import G from __main__ then:

--- t-pyannotate.py     (original)
+++ t-pyannotate.py     (refactored)
@@ -1,4 +1,5 @@
 from pyannotate_runtime import collect_types
+from __main__ import G

 collect_types.init_types_collection()
@@ -10,6 +11,7 @@

 def c(obj):
+    # type: (G) -> None
     pass

type_info.json:

[
    {
        "path": "t-pyannotate.py",
        "line": 8,
        "func_name": "G",
        "type_comments": [
            "() -> None"
        ],
        "samples": 1
    },
    {
        "path": "t-pyannotate.py",
        "line": 12,
        "func_name": "c",
        "type_comments": [
            "(__main__.G) -> None"
        ],
        "samples": 1
    }
]

pyannotate b7f96ca (current master).

[JukkaL]

This can sometimes be worked around by not running t-pyannotate.py directly and instead importing from another module that you run.