Commit | Line | Data |
---|---|---|
920dae64 AT |
1 | import sys |
2 | import os | |
3 | import marshal | |
4 | import imp | |
5 | import struct | |
6 | import time | |
7 | ||
8 | import zlib # implied prerequisite | |
9 | from zipfile import ZipFile, ZipInfo, ZIP_STORED, ZIP_DEFLATED | |
10 | from test import test_support | |
11 | from test.test_importhooks import ImportHooksBaseTestCase, test_src, test_co | |
12 | ||
13 | import zipimport | |
14 | ||
15 | ||
16 | def make_pyc(co, mtime): | |
17 | data = marshal.dumps(co) | |
18 | if type(mtime) is type(0.0): | |
19 | # Mac mtimes need a bit of special casing | |
20 | if mtime < 0x7fffffff: | |
21 | mtime = int(mtime) | |
22 | else: | |
23 | mtime = int(-0x100000000L + long(mtime)) | |
24 | pyc = imp.get_magic() + struct.pack("<i", int(mtime)) + data | |
25 | return pyc | |
26 | ||
27 | NOW = time.time() | |
28 | test_pyc = make_pyc(test_co, NOW) | |
29 | ||
30 | ||
31 | if __debug__: | |
32 | pyc_ext = ".pyc" | |
33 | else: | |
34 | pyc_ext = ".pyo" | |
35 | ||
36 | ||
37 | TESTMOD = "ziptestmodule" | |
38 | TESTPACK = "ziptestpackage" | |
39 | TESTPACK2 = "ziptestpackage2" | |
40 | TEMP_ZIP = os.path.abspath("junk95142" + os.extsep + "zip") | |
41 | ||
42 | class UncompressedZipImportTestCase(ImportHooksBaseTestCase): | |
43 | ||
44 | compression = ZIP_STORED | |
45 | ||
46 | def setUp(self): | |
47 | # We're reusing the zip archive path, so we must clear the | |
48 | # cached directory info. | |
49 | zipimport._zip_directory_cache.clear() | |
50 | ImportHooksBaseTestCase.setUp(self) | |
51 | ||
52 | def doTest(self, expected_ext, files, *modules, **kw): | |
53 | z = ZipFile(TEMP_ZIP, "w") | |
54 | try: | |
55 | for name, (mtime, data) in files.items(): | |
56 | zinfo = ZipInfo(name, time.localtime(mtime)) | |
57 | zinfo.compress_type = self.compression | |
58 | z.writestr(zinfo, data) | |
59 | z.close() | |
60 | ||
61 | stuff = kw.get("stuff", None) | |
62 | if stuff is not None: | |
63 | # Prepend 'stuff' to the start of the zipfile | |
64 | f = open(TEMP_ZIP, "rb") | |
65 | data = f.read() | |
66 | f.close() | |
67 | ||
68 | f = open(TEMP_ZIP, "wb") | |
69 | f.write(stuff) | |
70 | f.write(data) | |
71 | f.close() | |
72 | ||
73 | sys.path.insert(0, TEMP_ZIP) | |
74 | ||
75 | mod = __import__(".".join(modules), globals(), locals(), | |
76 | ["__dummy__"]) | |
77 | if expected_ext: | |
78 | file = mod.get_file() | |
79 | self.assertEquals(file, os.path.join(TEMP_ZIP, | |
80 | *modules) + expected_ext) | |
81 | finally: | |
82 | z.close() | |
83 | os.remove(TEMP_ZIP) | |
84 | ||
85 | def testAFakeZlib(self): | |
86 | # | |
87 | # This could cause a stack overflow before: importing zlib.py | |
88 | # from a compressed archive would cause zlib to be imported | |
89 | # which would find zlib.py in the archive, which would... etc. | |
90 | # | |
91 | # This test *must* be executed first: it must be the first one | |
92 | # to trigger zipimport to import zlib (zipimport caches the | |
93 | # zlib.decompress function object, after which the problem being | |
94 | # tested here wouldn't be a problem anymore... | |
95 | # (Hence the 'A' in the test method name: to make it the first | |
96 | # item in a list sorted by name, like unittest.makeSuite() does.) | |
97 | # | |
98 | # This test fails on platforms on which the zlib module is | |
99 | # statically linked, but the problem it tests for can't | |
100 | # occur in that case (builtin modules are always found first), | |
101 | # so we'll simply skip it then. Bug #765456. | |
102 | # | |
103 | if "zlib" in sys.builtin_module_names: | |
104 | return | |
105 | if "zlib" in sys.modules: | |
106 | del sys.modules["zlib"] | |
107 | files = {"zlib.py": (NOW, test_src)} | |
108 | try: | |
109 | self.doTest(".py", files, "zlib") | |
110 | except ImportError: | |
111 | if self.compression != ZIP_DEFLATED: | |
112 | self.fail("expected test to not raise ImportError") | |
113 | else: | |
114 | if self.compression != ZIP_STORED: | |
115 | self.fail("expected test to raise ImportError") | |
116 | ||
117 | def testPy(self): | |
118 | files = {TESTMOD + ".py": (NOW, test_src)} | |
119 | self.doTest(".py", files, TESTMOD) | |
120 | ||
121 | def testPyc(self): | |
122 | files = {TESTMOD + pyc_ext: (NOW, test_pyc)} | |
123 | self.doTest(pyc_ext, files, TESTMOD) | |
124 | ||
125 | def testBoth(self): | |
126 | files = {TESTMOD + ".py": (NOW, test_src), | |
127 | TESTMOD + pyc_ext: (NOW, test_pyc)} | |
128 | self.doTest(pyc_ext, files, TESTMOD) | |
129 | ||
130 | def testEmptyPy(self): | |
131 | files = {TESTMOD + ".py": (NOW, "")} | |
132 | self.doTest(None, files, TESTMOD) | |
133 | ||
134 | def testBadMagic(self): | |
135 | # make pyc magic word invalid, forcing loading from .py | |
136 | m0 = ord(test_pyc[0]) | |
137 | m0 ^= 0x04 # flip an arbitrary bit | |
138 | badmagic_pyc = chr(m0) + test_pyc[1:] | |
139 | files = {TESTMOD + ".py": (NOW, test_src), | |
140 | TESTMOD + pyc_ext: (NOW, badmagic_pyc)} | |
141 | self.doTest(".py", files, TESTMOD) | |
142 | ||
143 | def testBadMagic2(self): | |
144 | # make pyc magic word invalid, causing an ImportError | |
145 | m0 = ord(test_pyc[0]) | |
146 | m0 ^= 0x04 # flip an arbitrary bit | |
147 | badmagic_pyc = chr(m0) + test_pyc[1:] | |
148 | files = {TESTMOD + pyc_ext: (NOW, badmagic_pyc)} | |
149 | try: | |
150 | self.doTest(".py", files, TESTMOD) | |
151 | except ImportError: | |
152 | pass | |
153 | else: | |
154 | self.fail("expected ImportError; import from bad pyc") | |
155 | ||
156 | def testBadMTime(self): | |
157 | t3 = ord(test_pyc[7]) | |
158 | t3 ^= 0x02 # flip the second bit -- not the first as that one | |
159 | # isn't stored in the .py's mtime in the zip archive. | |
160 | badtime_pyc = test_pyc[:7] + chr(t3) + test_pyc[8:] | |
161 | files = {TESTMOD + ".py": (NOW, test_src), | |
162 | TESTMOD + pyc_ext: (NOW, badtime_pyc)} | |
163 | self.doTest(".py", files, TESTMOD) | |
164 | ||
165 | def testPackage(self): | |
166 | packdir = TESTPACK + os.sep | |
167 | files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc), | |
168 | packdir + TESTMOD + pyc_ext: (NOW, test_pyc)} | |
169 | self.doTest(pyc_ext, files, TESTPACK, TESTMOD) | |
170 | ||
171 | def testDeepPackage(self): | |
172 | packdir = TESTPACK + os.sep | |
173 | packdir2 = packdir + TESTPACK2 + os.sep | |
174 | files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc), | |
175 | packdir2 + "__init__" + pyc_ext: (NOW, test_pyc), | |
176 | packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)} | |
177 | self.doTest(pyc_ext, files, TESTPACK, TESTPACK2, TESTMOD) | |
178 | ||
179 | def testGetData(self): | |
180 | z = ZipFile(TEMP_ZIP, "w") | |
181 | z.compression = self.compression | |
182 | try: | |
183 | name = "testdata.dat" | |
184 | data = "".join([chr(x) for x in range(256)]) * 500 | |
185 | z.writestr(name, data) | |
186 | z.close() | |
187 | zi = zipimport.zipimporter(TEMP_ZIP) | |
188 | self.assertEquals(data, zi.get_data(name)) | |
189 | finally: | |
190 | z.close() | |
191 | os.remove(TEMP_ZIP) | |
192 | ||
193 | def testImporterAttr(self): | |
194 | src = """if 1: # indent hack | |
195 | def get_file(): | |
196 | return __file__ | |
197 | if __loader__.get_data("some.data") != "some data": | |
198 | raise AssertionError, "bad data"\n""" | |
199 | pyc = make_pyc(compile(src, "<???>", "exec"), NOW) | |
200 | files = {TESTMOD + pyc_ext: (NOW, pyc), | |
201 | "some.data": (NOW, "some data")} | |
202 | self.doTest(pyc_ext, files, TESTMOD) | |
203 | ||
204 | def testImport_WithStuff(self): | |
205 | # try importing from a zipfile which contains additional | |
206 | # stuff at the beginning of the file | |
207 | files = {TESTMOD + ".py": (NOW, test_src)} | |
208 | self.doTest(".py", files, TESTMOD, | |
209 | stuff="Some Stuff"*31) | |
210 | ||
211 | class CompressedZipImportTestCase(UncompressedZipImportTestCase): | |
212 | compression = ZIP_DEFLATED | |
213 | ||
214 | ||
215 | def test_main(): | |
216 | test_support.run_unittest( | |
217 | UncompressedZipImportTestCase, | |
218 | CompressedZipImportTestCase | |
219 | ) | |
220 | ||
221 | if __name__ == "__main__": | |
222 | test_main() |