rewrite _io imports to io in stubs (fixes #1) (PR #32)

Instagram · Dec 21, 2017 · d36c344 · d36c344
1 parent 7d2fafb
commit d36c344
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/monkeytype/stubs.py b/monkeytype/stubs.py
@@ -266,6 +266,8 @@ def render(self) -> str:
         imports = []
         for module in sorted(self.imports.keys()):
             names = sorted(self.imports[module])
+            if module == '_io':
+                module = module[1:]
             if len(names) == 1:
                 imports.append("from %s import %s" % (module, names[0]))
             else:

diff --git a/tests/test_stubs.py b/tests/test_stubs.py
@@ -74,6 +74,16 @@ def test_single_import(self):
         ])
         assert stub.render() == expected
 
+    def test_io_import_single(self):
+        """Single _io imports should convert to io"""
+        imports = ImportMap()
+        imports['_io'] = {'BytesIO'}
+        stub = ImportBlockStub(imports)
+        expected = "\n".join([
+            'from io import BytesIO',
+        ])
+        assert stub.render() == expected
+
     def test_multiple_imports(self):
         """Multiple imports from a single module should each be on their own line"""
         imports = ImportMap()
@@ -88,6 +98,19 @@ def test_multiple_imports(self):
         ])
         assert stub.render() == expected
 
+    def test_multiple_io_imports(self):
+        """Multiple imports from single _io module should be convert to io import"""
+        imports = ImportMap()
+        imports['_io'] = {'BytesIO', 'FileIO'}
+        stub = ImportBlockStub(imports)
+        expected = "\n".join([
+            'from io import (',
+            '    BytesIO,',
+            '    FileIO,',
+            ')',
+        ])
+        assert stub.render() == expected
+
 
 def simple_add(a: int, b: int) -> int:
     return a + b