Benjamin Peterson | ee8712c | 2008-05-20 21:35:26 +0000 | [diff] [blame] | 1 | from test.support import run_unittest |
Christian Heimes | dae2a89 | 2008-04-19 00:55:37 +0000 | [diff] [blame] | 2 | import unittest |
| 3 | import sys |
| 4 | import imp |
| 5 | import pkgutil |
| 6 | import os |
| 7 | import os.path |
| 8 | import tempfile |
| 9 | import shutil |
| 10 | import zipfile |
| 11 | |
| 12 | |
| 13 | |
| 14 | class PkgutilTests(unittest.TestCase): |
| 15 | |
| 16 | def setUp(self): |
| 17 | self.dirname = tempfile.mkdtemp() |
| 18 | sys.path.insert(0, self.dirname) |
| 19 | |
| 20 | def tearDown(self): |
| 21 | del sys.path[0] |
| 22 | shutil.rmtree(self.dirname) |
| 23 | |
| 24 | def test_getdata_filesys(self): |
| 25 | pkg = 'test_getdata_filesys' |
| 26 | |
| 27 | # Include a LF and a CRLF, to test that binary data is read back |
| 28 | RESOURCE_DATA = b'Hello, world!\nSecond line\r\nThird line' |
| 29 | |
| 30 | # Make a package with some resources |
| 31 | package_dir = os.path.join(self.dirname, pkg) |
| 32 | os.mkdir(package_dir) |
| 33 | # Empty init.py |
| 34 | f = open(os.path.join(package_dir, '__init__.py'), "wb") |
| 35 | f.close() |
| 36 | # Resource files, res.txt, sub/res.txt |
| 37 | f = open(os.path.join(package_dir, 'res.txt'), "wb") |
| 38 | f.write(RESOURCE_DATA) |
| 39 | f.close() |
| 40 | os.mkdir(os.path.join(package_dir, 'sub')) |
| 41 | f = open(os.path.join(package_dir, 'sub', 'res.txt'), "wb") |
| 42 | f.write(RESOURCE_DATA) |
| 43 | f.close() |
| 44 | |
| 45 | # Check we can read the resources |
| 46 | res1 = pkgutil.get_data(pkg, 'res.txt') |
| 47 | self.assertEqual(res1, RESOURCE_DATA) |
| 48 | res2 = pkgutil.get_data(pkg, 'sub/res.txt') |
| 49 | self.assertEqual(res2, RESOURCE_DATA) |
| 50 | |
| 51 | del sys.modules[pkg] |
| 52 | |
| 53 | def test_getdata_zipfile(self): |
| 54 | zip = 'test_getdata_zipfile.zip' |
| 55 | pkg = 'test_getdata_zipfile' |
| 56 | |
| 57 | # Include a LF and a CRLF, to test that binary data is read back |
| 58 | RESOURCE_DATA = b'Hello, world!\nSecond line\r\nThird line' |
| 59 | |
| 60 | # Make a package with some resources |
| 61 | zip_file = os.path.join(self.dirname, zip) |
| 62 | z = zipfile.ZipFile(zip_file, 'w') |
| 63 | |
| 64 | # Empty init.py |
| 65 | z.writestr(pkg + '/__init__.py', "") |
| 66 | # Resource files, res.txt, sub/res.txt |
| 67 | z.writestr(pkg + '/res.txt', RESOURCE_DATA) |
| 68 | z.writestr(pkg + '/sub/res.txt', RESOURCE_DATA) |
| 69 | z.close() |
| 70 | |
| 71 | # Check we can read the resources |
| 72 | sys.path.insert(0, zip_file) |
| 73 | res1 = pkgutil.get_data(pkg, 'res.txt') |
| 74 | self.assertEqual(res1, RESOURCE_DATA) |
| 75 | res2 = pkgutil.get_data(pkg, 'sub/res.txt') |
| 76 | self.assertEqual(res2, RESOURCE_DATA) |
Georg Brandl | ae83d6e | 2009-08-13 09:04:31 +0000 | [diff] [blame] | 77 | |
| 78 | names = [] |
| 79 | for loader, name, ispkg in pkgutil.iter_modules([zip_file]): |
| 80 | names.append(name) |
| 81 | self.assertEqual(names, ['test_getdata_zipfile']) |
| 82 | |
Christian Heimes | dae2a89 | 2008-04-19 00:55:37 +0000 | [diff] [blame] | 83 | del sys.path[0] |
| 84 | |
| 85 | del sys.modules[pkg] |
| 86 | |
| 87 | class PkgutilPEP302Tests(unittest.TestCase): |
| 88 | |
| 89 | class MyTestLoader(object): |
| 90 | def load_module(self, fullname): |
| 91 | # Create an empty module |
| 92 | mod = sys.modules.setdefault(fullname, imp.new_module(fullname)) |
| 93 | mod.__file__ = "<%s>" % self.__class__.__name__ |
| 94 | mod.__loader__ = self |
| 95 | # Make it a package |
| 96 | mod.__path__ = [] |
| 97 | # Count how many times the module is reloaded |
| 98 | mod.__dict__['loads'] = mod.__dict__.get('loads',0) + 1 |
| 99 | return mod |
| 100 | |
| 101 | def get_data(self, path): |
| 102 | return "Hello, world!" |
| 103 | |
| 104 | class MyTestImporter(object): |
| 105 | def find_module(self, fullname, path=None): |
| 106 | return PkgutilPEP302Tests.MyTestLoader() |
| 107 | |
| 108 | def setUp(self): |
| 109 | sys.meta_path.insert(0, self.MyTestImporter()) |
| 110 | |
| 111 | def tearDown(self): |
| 112 | del sys.meta_path[0] |
| 113 | |
| 114 | def test_getdata_pep302(self): |
| 115 | # Use a dummy importer/loader |
| 116 | self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!") |
| 117 | del sys.modules['foo'] |
| 118 | |
| 119 | def test_alreadyloaded(self): |
| 120 | # Ensure that get_data works without reloading - the "loads" module |
| 121 | # variable in the example loader should count how many times a reload |
| 122 | # occurs. |
| 123 | import foo |
| 124 | self.assertEqual(foo.loads, 1) |
| 125 | self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!") |
| 126 | self.assertEqual(foo.loads, 1) |
| 127 | del sys.modules['foo'] |
| 128 | |
| 129 | def test_main(): |
| 130 | run_unittest(PkgutilTests, PkgutilPEP302Tests) |
Christian Heimes | e57950f | 2008-04-21 13:08:03 +0000 | [diff] [blame] | 131 | # this is necessary if test is run repeated (like when finding leaks) |
| 132 | import zipimport |
| 133 | zipimport._zip_directory_cache.clear() |
Christian Heimes | dae2a89 | 2008-04-19 00:55:37 +0000 | [diff] [blame] | 134 | |
| 135 | if __name__ == '__main__': |
| 136 | test_main() |