14adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaofrom test.test_support import run_unittest
24adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport unittest
34adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport sys
44adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport imp
54adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport pkgutil
64adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport os
74adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport os.path
84adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport tempfile
94adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport shutil
104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoimport zipfile
114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass PkgutilTests(unittest.TestCase):
154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def setUp(self):
174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.dirname = tempfile.mkdtemp()
184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.addCleanup(shutil.rmtree, self.dirname)
194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        sys.path.insert(0, self.dirname)
204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def tearDown(self):
224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.path[0]
234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def test_getdata_filesys(self):
254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        pkg = 'test_getdata_filesys'
264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Include a LF and a CRLF, to test that binary data is read back
284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        RESOURCE_DATA = 'Hello, world!\nSecond line\r\nThird line'
294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Make a package with some resources
314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        package_dir = os.path.join(self.dirname, pkg)
324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        os.mkdir(package_dir)
334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Empty init.py
344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f = open(os.path.join(package_dir, '__init__.py'), "wb")
354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f.close()
364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Resource files, res.txt, sub/res.txt
374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f = open(os.path.join(package_dir, 'res.txt'), "wb")
384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f.write(RESOURCE_DATA)
394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f.close()
404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        os.mkdir(os.path.join(package_dir, 'sub'))
414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f = open(os.path.join(package_dir, 'sub', 'res.txt'), "wb")
424adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f.write(RESOURCE_DATA)
434adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        f.close()
444adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
454adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Check we can read the resources
464adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        res1 = pkgutil.get_data(pkg, 'res.txt')
474adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(res1, RESOURCE_DATA)
484adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        res2 = pkgutil.get_data(pkg, 'sub/res.txt')
494adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(res2, RESOURCE_DATA)
504adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
514adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.modules[pkg]
524adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
534adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def test_getdata_zipfile(self):
544adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        zip = 'test_getdata_zipfile.zip'
554adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        pkg = 'test_getdata_zipfile'
564adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
574adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Include a LF and a CRLF, to test that binary data is read back
584adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        RESOURCE_DATA = 'Hello, world!\nSecond line\r\nThird line'
594adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
604adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Make a package with some resources
614adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        zip_file = os.path.join(self.dirname, zip)
624adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        z = zipfile.ZipFile(zip_file, 'w')
634adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
644adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Empty init.py
654adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        z.writestr(pkg + '/__init__.py', "")
664adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Resource files, res.txt, sub/res.txt
674adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        z.writestr(pkg + '/res.txt', RESOURCE_DATA)
684adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        z.writestr(pkg + '/sub/res.txt', RESOURCE_DATA)
694adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        z.close()
704adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
714adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Check we can read the resources
724adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        sys.path.insert(0, zip_file)
734adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        res1 = pkgutil.get_data(pkg, 'res.txt')
744adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(res1, RESOURCE_DATA)
754adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        res2 = pkgutil.get_data(pkg, 'sub/res.txt')
764adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(res2, RESOURCE_DATA)
774adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.path[0]
784adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
794adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.modules[pkg]
804adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
814adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def test_unreadable_dir_on_syspath(self):
824adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # issue7367 - walk_packages failed if unreadable dir on sys.path
834adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        package_name = "unreadable_package"
844adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        d = os.path.join(self.dirname, package_name)
854adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # this does not appear to create an unreadable dir on Windows
864adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        #   but the test should not fail anyway
874adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        os.mkdir(d, 0)
884adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.addCleanup(os.rmdir, d)
894adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        for t in pkgutil.walk_packages(path=[self.dirname]):
904adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            self.fail("unexpected package found")
914adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
924adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoclass PkgutilPEP302Tests(unittest.TestCase):
934adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
944adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    class MyTestLoader(object):
954adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        def load_module(self, fullname):
964adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            # Create an empty module
974adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
984adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            mod.__file__ = "<%s>" % self.__class__.__name__
994adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            mod.__loader__ = self
1004adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            # Make it a package
1014adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            mod.__path__ = []
1024adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            # Count how many times the module is reloaded
1034adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            mod.__dict__['loads'] = mod.__dict__.get('loads',0) + 1
1044adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            return mod
1054adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1064adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        def get_data(self, path):
1074adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            return "Hello, world!"
1084adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1094adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    class MyTestImporter(object):
1104adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        def find_module(self, fullname, path=None):
1114adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao            return PkgutilPEP302Tests.MyTestLoader()
1124adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1134adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def setUp(self):
1144adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        sys.meta_path.insert(0, self.MyTestImporter())
1154adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1164adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def tearDown(self):
1174adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.meta_path[0]
1184adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1194adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def test_getdata_pep302(self):
1204adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Use a dummy importer/loader
1214adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
1224adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.modules['foo']
1234adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1244adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    def test_alreadyloaded(self):
1254adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # Ensure that get_data works without reloading - the "loads" module
1264adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # variable in the example loader should count how many times a reload
1274adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        # occurs.
1284adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        import foo
1294adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(foo.loads, 1)
1304adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
1314adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        self.assertEqual(foo.loads, 1)
1324adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao        del sys.modules['foo']
1334adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1344adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaodef test_main():
1354adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    run_unittest(PkgutilTests, PkgutilPEP302Tests)
1364adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    # this is necessary if test is run repeated (like when finding leaks)
1374adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    import zipimport
1384adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    zipimport._zip_directory_cache.clear()
1394adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao
1404adfde8bc82dd39f59e0445588c3e599ada477dJosh Gaoif __name__ == '__main__':
1414adfde8bc82dd39f59e0445588c3e599ada477dJosh Gao    test_main()
142