Refactor testpep561 and add test (#5782)

Adds tests for #5785. Closes #5767.

mypy/test/testpep561.py

@@ -1,4 +1,5 @@
 from contextlib import contextmanager
+from enum import Enum
 import os
 import sys
 import tempfile
@@ -21,38 +22,105 @@
 reveal_type(a)
 """
 
-NAMESPACE_PROGRAM = """
-from typedpkg_nested.nested_package.nested_module import nested_func
-from typedpkg_namespace.alpha.alpha_module import alpha_func
+_NAMESPACE_PROGRAM = """
+{import_style}
+from typedpkg_ns.ns.dne import dne
 
-nested_func("abc")
-alpha_func(False)
+af("abc")
+bf(False)
+dne(123)
 
-nested_func(False)
-alpha_func(2)
+af(False)
+bf(2)
+dne("abc")
 """
 
 
-def check_mypy_run(cmd_line: List[str],
-                   python_executable: str = sys.executable,
-                   expected_out: str = '',
-                   expected_err: str = '',
-                   expected_returncode: int = 1,
-                   venv_dir: Optional[str] = None) -> None:
-    """Helper to run mypy and check the output."""
-    if venv_dir is not None:
-        old_dir = os.getcwd()
-        os.chdir(venv_dir)
-    try:
-        if python_executable != sys.executable:
-            cmd_line.append('--python-executable={}'.format(python_executable))
-        out, err, returncode = mypy.api.run(cmd_line)
-        assert out == expected_out, err
-        assert err == expected_err, out
-        assert returncode == expected_returncode, returncode
-    finally:
+class NSImportStyle(Enum):
+    # These should all be on exactly two lines because NamespaceMsg
+    # uses line numbers which expect the imports to be exactly two lines
+    from_import = """\
+from typedpkg.pkg.aaa import af
+from typedpkg_ns.ns.bbb import bf"""
+    import_as = """\
+import typedpkg.pkg.aaa as nm; af = nm.af
+import typedpkg_ns.ns.bbb as am; bf = am.bf"""
+    reg_import = """\
+import typedpkg.pkg.aaa; af = typedpkg.pkg.aaa.af
+import typedpkg_ns.ns.bbb; bf = typedpkg_ns.ns.bbb.bf"""
+
+
+class SimpleMsg(Enum):
+    msg_dne = "{tempfile}:3: error: Module 'typedpkg' has no attribute 'dne'"
+    msg_list = "{tempfile}:5: error: Revealed type is 'builtins.list[builtins.str]'"
+    msg_tuple = "{tempfile}:5: error: Revealed type is 'builtins.tuple[builtins.str]'"
+
+
+class NamespaceMsg(Enum):
+    cfm_beta = ("{tempfile}:4: error: Cannot find module named "
+                "'typedpkg_ns.ns.dne'")
+    help_note = ('{tempfile}:4: note: (Perhaps setting MYPYPATH or using the '
+                 '"--ignore-missing-imports" flag would help)')
+    bool_str = ('{tempfile}:10: error: Argument 1 has incompatible type '
+                '"bool"; expected "str"')
+    int_bool = ('{tempfile}:11: error: Argument 1 has incompatible type '
+                '"int"; expected "bool"')
+    to_bool_str = ('{tempfile}:10: error: Argument 1 to "af" has incompatible type '
+                   '"bool"; expected "str"')
+    to_int_bool = ('{tempfile}:11: error: Argument 1 to "bf" has incompatible type '
+                   '"int"; expected "bool"')
+
+
+def create_ns_program_src(import_style: NSImportStyle) -> str:
+    return _NAMESPACE_PROGRAM.format(import_style=import_style.value)
+
+
+class ExampleProg(object):
+    _fname = 'test_program.py'
+
+    def __init__(self, source_code: str) -> None:
+        self._source_code = source_code
+
+        self._temp_dir = None  # type: Optional[tempfile.TemporaryDirectory[str]]
+        self._full_fname = ''
+
+    def create(self) -> None:
+        self._temp_dir = tempfile.TemporaryDirectory()
+        self._full_fname = os.path.join(self._temp_dir.name, self._fname)
+        with open(self._full_fname, 'w+') as f:
+            f.write(self._source_code)
+
+    def cleanup(self) -> None:
+        if self._temp_dir:
+            self._temp_dir.cleanup()
+
+    def build_msg(self, *msgs: Enum) -> str:
+        return '\n'.join(
+            msg.value.format(tempfile=self._full_fname)
+            for msg in msgs
+        ) + '\n'
+
+    def check_mypy_run(self,
+                       python_executable: str,
+                       expected_out: List[Enum],
+                       expected_err: str = '',
+                       expected_returncode: int = 1,
+                       venv_dir: Optional[str] = None) -> None:
+        """Helper to run mypy and check the output."""
+        cmd_line = [self._full_fname]
         if venv_dir is not None:
-            os.chdir(old_dir)
+            old_dir = os.getcwd()
+            os.chdir(venv_dir)
+        try:
+            if python_executable != sys.executable:
+                cmd_line.append('--python-executable={}'.format(python_executable))
+            out, err, returncode = mypy.api.run(cmd_line)
+            assert out == self.build_msg(*expected_out), err
+            assert err == expected_err, out
+            assert returncode == expected_returncode, returncode
+        finally:
+            if venv_dir is not None:
+                os.chdir(old_dir)
 
 
 class TestPEP561(TestCase):
@@ -102,127 +170,157 @@ def install_package(self, pkg: str,
             self.fail('\n'.join(lines))
 
     def setUp(self) -> None:
-        self.temp_file_dir = tempfile.TemporaryDirectory()
-        self.tempfile = os.path.join(self.temp_file_dir.name, 'simple.py')
-        with open(self.tempfile, 'w+') as file:
-            file.write(SIMPLE_PROGRAM)
-        self.namespace_tempfile = os.path.join(self.temp_file_dir.name, 'namespace_program.py')
-        with open(self.namespace_tempfile, 'w+') as file:
-            file.write(NAMESPACE_PROGRAM)
-
-        self.msg_dne = \
-            "{}:3: error: Module 'typedpkg' has no attribute 'dne'\n".format(self.tempfile)
-        self.msg_list = \
-            "{}:5: error: Revealed type is 'builtins.list[builtins.str]'\n".format(self.tempfile)
-        self.msg_tuple = \
-            "{}:5: error: Revealed type is 'builtins.tuple[builtins.str]'\n".format(self.tempfile)
-
-        self.namespace_msg_bool_str = (
-            '{0}:8: error: Argument 1 to "nested_func" has incompatible type "bool"; '
-            'expected "str"\n'.format(self.namespace_tempfile))
-        self.namespace_msg_int_bool = (
-            '{0}:9: error: Argument 1 to "alpha_func" has incompatible type "int"; '
-            'expected "bool"\n'.format(self.namespace_tempfile))
+        self.simple_prog = ExampleProg(SIMPLE_PROGRAM)
+        self.from_ns_prog = ExampleProg(create_ns_program_src(NSImportStyle.from_import))
+        self.import_as_ns_prog = ExampleProg(create_ns_program_src(NSImportStyle.import_as))
+        self.regular_import_ns_prog = ExampleProg(create_ns_program_src(NSImportStyle.reg_import))
 
     def tearDown(self) -> None:
-        self.temp_file_dir.cleanup()
+        self.simple_prog.cleanup()
+        self.from_ns_prog.cleanup()
+        self.import_as_ns_prog.cleanup()
+        self.regular_import_ns_prog.cleanup()
 
     def test_get_pkg_dirs(self) -> None:
         """Check that get_package_dirs works."""
         dirs = get_site_packages_dirs(sys.executable)
         assert dirs
 
     def test_typedpkg_stub_package(self) -> None:
+        self.simple_prog.create()
         with self.virtualenv() as venv:
             venv_dir, python_executable = venv
             self.install_package('typedpkg-stubs', python_executable)
-            check_mypy_run(
-                [self.tempfile],
+            self.simple_prog.check_mypy_run(
                 python_executable,
-                expected_out=self.msg_dne + self.msg_list,
+                [SimpleMsg.msg_dne, SimpleMsg.msg_list],
                 venv_dir=venv_dir,
             )
 
     def test_typedpkg(self) -> None:
+        self.simple_prog.create()
         with self.virtualenv() as venv:
             venv_dir, python_executable = venv
             self.install_package('typedpkg', python_executable)
-            check_mypy_run(
-                [self.tempfile],
+            self.simple_prog.check_mypy_run(
                 python_executable,
-                expected_out=self.msg_tuple,
+                [SimpleMsg.msg_tuple],
                 venv_dir=venv_dir,
             )
 
     def test_stub_and_typed_pkg(self) -> None:
+        self.simple_prog.create()
         with self.virtualenv() as venv:
             venv_dir, python_executable = venv
             self.install_package('typedpkg', python_executable)
             self.install_package('typedpkg-stubs', python_executable)
-            check_mypy_run(
-                [self.tempfile],
+            self.simple_prog.check_mypy_run(
                 python_executable,
-                expected_out=self.msg_list,
+                [SimpleMsg.msg_list],
                 venv_dir=venv_dir,
             )
 
     def test_typedpkg_stubs_python2(self) -> None:
+        self.simple_prog.create()
         python2 = try_find_python2_interpreter()
         if python2:
             with self.virtualenv(python2) as venv:
                 venv_dir, py2 = venv
                 self.install_package('typedpkg-stubs', py2)
-                check_mypy_run(
-                    [self.tempfile],
+                self.simple_prog.check_mypy_run(
                     py2,
-                    expected_out=self.msg_dne + self.msg_list,
+                    [SimpleMsg.msg_dne, SimpleMsg.msg_list],
                     venv_dir=venv_dir,
                 )
 
     def test_typedpkg_python2(self) -> None:
+        self.simple_prog.create()
         python2 = try_find_python2_interpreter()
         if python2:
             with self.virtualenv(python2) as venv:
                 venv_dir, py2 = venv
                 self.install_package('typedpkg', py2)
-                check_mypy_run(
-                    [self.tempfile],
+                self.simple_prog.check_mypy_run(
                     py2,
-                    expected_out=self.msg_tuple,
+                    [SimpleMsg.msg_tuple],
                     venv_dir=venv_dir,
                 )
 
     def test_typedpkg_egg(self) -> None:
+        self.simple_prog.create()
         with self.virtualenv() as venv:
             venv_dir, python_executable = venv
             self.install_package('typedpkg', python_executable, use_pip=False)
-            check_mypy_run(
-                [self.tempfile],
+            self.simple_prog.check_mypy_run(
                 python_executable,
-                expected_out=self.msg_tuple,
+                [SimpleMsg.msg_tuple],
                 venv_dir=venv_dir,
             )
 
     def test_typedpkg_editable(self) -> None:
+        self.simple_prog.create()
         with self.virtualenv() as venv:
             venv_dir, python_executable = venv
             self.install_package('typedpkg', python_executable, editable=True)
-            check_mypy_run(
-                [self.tempfile],
+            self.simple_prog.check_mypy_run(
                 python_executable,
-                expected_out=self.msg_tuple,
+                [SimpleMsg.msg_tuple],
                 venv_dir=venv_dir,
             )
 
-    def test_nested_and_namespace(self) -> None:
+    def test_typedpkg_egg_editable(self) -> None:
+        self.simple_prog.create()
         with self.virtualenv() as venv:
             venv_dir, python_executable = venv
-            self.install_package('typedpkg_nested', python_executable)
-            self.install_package('typedpkg_namespace-alpha', python_executable)
-            check_mypy_run(
-                [self.namespace_tempfile],
+            self.install_package('typedpkg', python_executable, use_pip=False, editable=True)
+            self.simple_prog.check_mypy_run(
+                python_executable,
+                [SimpleMsg.msg_tuple],
+                venv_dir=venv_dir,
+            )
+
+    def test_nested_and_namespace_from_import(self) -> None:
+        self.from_ns_prog.create()
+        with self.virtualenv() as venv:
+            venv_dir, python_executable = venv
+            self.install_package('typedpkg', python_executable)
+            self.install_package('typedpkg_ns', python_executable)
+            self.from_ns_prog.check_mypy_run(
+                python_executable,
+                [NamespaceMsg.cfm_beta,
+                 NamespaceMsg.help_note,
+                 NamespaceMsg.to_bool_str,
+                 NamespaceMsg.to_int_bool],
+                venv_dir=venv_dir,
+            )
+
+    def test_nested_and_namespace_import_as(self) -> None:
+        self.import_as_ns_prog.create()
+        with self.virtualenv() as venv:
+            venv_dir, python_executable = venv
+            self.install_package('typedpkg', python_executable)
+            self.install_package('typedpkg_ns', python_executable)
+            self.import_as_ns_prog.check_mypy_run(
+                python_executable,
+                [NamespaceMsg.cfm_beta,
+                 NamespaceMsg.help_note,
+                 NamespaceMsg.bool_str,
+                 NamespaceMsg.int_bool],
+                venv_dir=venv_dir,
+            )
+
+    def test_nested_and_namespace_regular_import(self) -> None:
+        self.regular_import_ns_prog.create()
+        with self.virtualenv() as venv:
+            venv_dir, python_executable = venv
+            self.install_package('typedpkg', python_executable)
+            self.install_package('typedpkg_ns', python_executable)
+            self.regular_import_ns_prog.check_mypy_run(
                 python_executable,
-                expected_out=self.namespace_msg_bool_str + self.namespace_msg_int_bool,
+                [NamespaceMsg.cfm_beta,
+                 NamespaceMsg.help_note,
+                 NamespaceMsg.bool_str,
+                 NamespaceMsg.int_bool],
                 venv_dir=venv_dir,
             )
 

test-data/packages/typedpkg/setup.py

@@ -9,7 +9,7 @@
     author="The mypy team",
     version='0.1',
     package_data={'typedpkg': ['py.typed']},
-    packages=['typedpkg'],
+    packages=['typedpkg', 'typedpkg.pkg'],
     include_package_data=True,
     zip_safe=False,
 )

test-data/packages/typedpkg/typedpkg/pkg/aaa.py

@@ -0,0 +1,2 @@
+def af(a: str) -> str:
+    return a + " nested"

test-data/packages/typedpkg_namespace-alpha/setup.py → test-data/packages/typedpkg_ns/setup.py

@@ -4,7 +4,7 @@
     name='typedpkg_namespace.alpha',
     version='1.0.0',
     packages=find_packages(),
-    namespace_packages=['typedpkg_namespace'],
+    namespace_packages=['typedpkg_ns'],
     zip_safe=False,
-    package_data={'typedpkg_namespace.alpha': ['py.typed']}
+    package_data={'typedpkg_ns.ns': ['py.typed']}
 )

test-data/packages/typedpkg_ns/typedpkg_ns/ns/bbb.py

@@ -0,0 +1,2 @@
+def bf(a: bool) -> bool:
+    return not a