1from test.test_support import run_unittest
2import unittest
3import sys
4import imp
5import pkgutil
6import os
7import os.path
8import tempfile
9import shutil
10import zipfile
11
12
13
14class PkgutilTests(unittest.TestCase):
15
16    def setUp(self):
17        self.dirname = tempfile.mkdtemp()
18        self.addCleanup(shutil.rmtree, self.dirname)
19        sys.path.insert(0, self.dirname)
20
21    def tearDown(self):
22        del sys.path[0]
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 = '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 = '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)
77        del sys.path[0]
78
79        del sys.modules[pkg]
80
81    def test_unreadable_dir_on_syspath(self):
82        # issue7367 - walk_packages failed if unreadable dir on sys.path
83        package_name = "unreadable_package"
84        d = os.path.join(self.dirname, package_name)
85        # this does not appear to create an unreadable dir on Windows
86        #   but the test should not fail anyway
87        os.mkdir(d, 0)
88        self.addCleanup(os.rmdir, d)
89        for t in pkgutil.walk_packages(path=[self.dirname]):
90            self.fail("unexpected package found")
91
92class PkgutilPEP302Tests(unittest.TestCase):
93
94    class MyTestLoader(object):
95        def load_module(self, fullname):
96            # Create an empty module
97            mod = sys.modules.setdefault(fullname, imp.new_module(fullname))
98            mod.__file__ = "<%s>" % self.__class__.__name__
99            mod.__loader__ = self
100            # Make it a package
101            mod.__path__ = []
102            # Count how many times the module is reloaded
103            mod.__dict__['loads'] = mod.__dict__.get('loads',0) + 1
104            return mod
105
106        def get_data(self, path):
107            return "Hello, world!"
108
109    class MyTestImporter(object):
110        def find_module(self, fullname, path=None):
111            return PkgutilPEP302Tests.MyTestLoader()
112
113    def setUp(self):
114        sys.meta_path.insert(0, self.MyTestImporter())
115
116    def tearDown(self):
117        del sys.meta_path[0]
118
119    def test_getdata_pep302(self):
120        # Use a dummy importer/loader
121        self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
122        del sys.modules['foo']
123
124    def test_alreadyloaded(self):
125        # Ensure that get_data works without reloading - the "loads" module
126        # variable in the example loader should count how many times a reload
127        # occurs.
128        import foo
129        self.assertEqual(foo.loads, 1)
130        self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")
131        self.assertEqual(foo.loads, 1)
132        del sys.modules['foo']
133
134def test_main():
135    run_unittest(PkgutilTests, PkgutilPEP302Tests)
136    # this is necessary if test is run repeated (like when finding leaks)
137    import zipimport
138    zipimport._zip_directory_cache.clear()
139
140if __name__ == '__main__':
141    test_main()
142