/Lib/test/test_zipimport.py

http://unladen-swallow.googlecode.com/ · Python · 482 lines · 369 code · 78 blank · 35 comment · 44 complexity · 7cf6753c6fd404c1c6071e8cf620b231 MD5 · raw file

  1. import sys
  2. import os
  3. import marshal
  4. import imp
  5. import struct
  6. import time
  7. import unittest
  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. import zipimport
  13. import linecache
  14. import doctest
  15. import inspect
  16. import StringIO
  17. from traceback import extract_tb, extract_stack, print_tb
  18. raise_src = 'def do_raise(): raise TypeError\n'
  19. # so we only run testAFakeZlib once if this test is run repeatedly
  20. # which happens when we look for ref leaks
  21. test_imported = False
  22. def make_pyc(co, mtime):
  23. data = marshal.dumps(co)
  24. if type(mtime) is type(0.0):
  25. # Mac mtimes need a bit of special casing
  26. if mtime < 0x7fffffff:
  27. mtime = int(mtime)
  28. else:
  29. mtime = int(-0x100000000L + long(mtime))
  30. pyc = imp.get_magic() + struct.pack("<i", int(mtime)) + data
  31. return pyc
  32. def module_path_to_dotted_name(path):
  33. return path.replace(os.sep, '.')
  34. NOW = time.time()
  35. test_pyc = make_pyc(test_co, NOW)
  36. if __debug__:
  37. pyc_ext = ".pyc"
  38. else:
  39. pyc_ext = ".pyo"
  40. TESTMOD = "ziptestmodule"
  41. TESTPACK = "ziptestpackage"
  42. TESTPACK2 = "ziptestpackage2"
  43. TEMP_ZIP = os.path.abspath("junk95142" + os.extsep + "zip")
  44. class UncompressedZipImportTestCase(ImportHooksBaseTestCase):
  45. compression = ZIP_STORED
  46. def setUp(self):
  47. # We're reusing the zip archive path, so we must clear the
  48. # cached directory info and linecache
  49. linecache.clearcache()
  50. zipimport._zip_directory_cache.clear()
  51. ImportHooksBaseTestCase.setUp(self)
  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. stuff = kw.get("stuff", None)
  61. if stuff is not None:
  62. # Prepend 'stuff' to the start of the zipfile
  63. f = open(TEMP_ZIP, "rb")
  64. data = f.read()
  65. f.close()
  66. f = open(TEMP_ZIP, "wb")
  67. f.write(stuff)
  68. f.write(data)
  69. f.close()
  70. sys.path.insert(0, TEMP_ZIP)
  71. mod = __import__(".".join(modules), globals(), locals(),
  72. ["__dummy__"])
  73. call = kw.get('call')
  74. if call is not None:
  75. call(mod)
  76. if expected_ext:
  77. file = mod.get_file()
  78. self.assertEquals(file, os.path.join(TEMP_ZIP,
  79. *modules) + expected_ext)
  80. finally:
  81. z.close()
  82. os.remove(TEMP_ZIP)
  83. def testAFakeZlib(self):
  84. #
  85. # This could cause a stack overflow before: importing zlib.py
  86. # from a compressed archive would cause zlib to be imported
  87. # which would find zlib.py in the archive, which would... etc.
  88. #
  89. # This test *must* be executed first: it must be the first one
  90. # to trigger zipimport to import zlib (zipimport caches the
  91. # zlib.decompress function object, after which the problem being
  92. # tested here wouldn't be a problem anymore...
  93. # (Hence the 'A' in the test method name: to make it the first
  94. # item in a list sorted by name, like unittest.makeSuite() does.)
  95. #
  96. # This test fails on platforms on which the zlib module is
  97. # statically linked, but the problem it tests for can't
  98. # occur in that case (builtin modules are always found first),
  99. # so we'll simply skip it then. Bug #765456.
  100. #
  101. if "zlib" in sys.builtin_module_names:
  102. return
  103. if "zlib" in sys.modules:
  104. del sys.modules["zlib"]
  105. files = {"zlib.py": (NOW, test_src)}
  106. try:
  107. self.doTest(".py", files, "zlib")
  108. except ImportError:
  109. if self.compression != ZIP_DEFLATED:
  110. self.fail("expected test to not raise ImportError")
  111. else:
  112. if self.compression != ZIP_STORED:
  113. self.fail("expected test to raise ImportError")
  114. def testPy(self):
  115. files = {TESTMOD + ".py": (NOW, test_src)}
  116. self.doTest(".py", files, TESTMOD)
  117. def testPyc(self):
  118. files = {TESTMOD + pyc_ext: (NOW, test_pyc)}
  119. self.doTest(pyc_ext, files, TESTMOD)
  120. def testBoth(self):
  121. files = {TESTMOD + ".py": (NOW, test_src),
  122. TESTMOD + pyc_ext: (NOW, test_pyc)}
  123. self.doTest(pyc_ext, files, TESTMOD)
  124. def testEmptyPy(self):
  125. files = {TESTMOD + ".py": (NOW, "")}
  126. self.doTest(None, files, TESTMOD)
  127. def testBadMagic(self):
  128. # make pyc magic word invalid, forcing loading from .py
  129. m0 = ord(test_pyc[0])
  130. m0 ^= 0x04 # flip an arbitrary bit
  131. badmagic_pyc = chr(m0) + test_pyc[1:]
  132. files = {TESTMOD + ".py": (NOW, test_src),
  133. TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
  134. self.doTest(".py", files, TESTMOD)
  135. def testBadMagic2(self):
  136. # make pyc magic word invalid, causing an ImportError
  137. m0 = ord(test_pyc[0])
  138. m0 ^= 0x04 # flip an arbitrary bit
  139. badmagic_pyc = chr(m0) + test_pyc[1:]
  140. files = {TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
  141. try:
  142. self.doTest(".py", files, TESTMOD)
  143. except ImportError:
  144. pass
  145. else:
  146. self.fail("expected ImportError; import from bad pyc")
  147. def testBadMTime(self):
  148. t3 = ord(test_pyc[7])
  149. t3 ^= 0x02 # flip the second bit -- not the first as that one
  150. # isn't stored in the .py's mtime in the zip archive.
  151. badtime_pyc = test_pyc[:7] + chr(t3) + test_pyc[8:]
  152. files = {TESTMOD + ".py": (NOW, test_src),
  153. TESTMOD + pyc_ext: (NOW, badtime_pyc)}
  154. self.doTest(".py", files, TESTMOD)
  155. def testPackage(self):
  156. packdir = TESTPACK + os.sep
  157. files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
  158. packdir + TESTMOD + pyc_ext: (NOW, test_pyc)}
  159. self.doTest(pyc_ext, files, TESTPACK, TESTMOD)
  160. def testDeepPackage(self):
  161. packdir = TESTPACK + os.sep
  162. packdir2 = packdir + TESTPACK2 + os.sep
  163. files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
  164. packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
  165. packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
  166. self.doTest(pyc_ext, files, TESTPACK, TESTPACK2, TESTMOD)
  167. def testZipImporterMethods(self):
  168. packdir = TESTPACK + os.sep
  169. packdir2 = packdir + TESTPACK2 + os.sep
  170. files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
  171. packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
  172. packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
  173. z = ZipFile(TEMP_ZIP, "w")
  174. try:
  175. for name, (mtime, data) in files.items():
  176. zinfo = ZipInfo(name, time.localtime(mtime))
  177. zinfo.compress_type = self.compression
  178. z.writestr(zinfo, data)
  179. z.close()
  180. zi = zipimport.zipimporter(TEMP_ZIP)
  181. self.assertEquals(zi.archive, TEMP_ZIP)
  182. self.assertEquals(zi.is_package(TESTPACK), True)
  183. mod = zi.load_module(TESTPACK)
  184. self.assertEquals(zi._get_filename(TESTPACK), mod.__file__)
  185. self.assertEquals(zi.is_package(packdir + '__init__'), False)
  186. self.assertEquals(zi.is_package(packdir + TESTPACK2), True)
  187. self.assertEquals(zi.is_package(packdir2 + TESTMOD), False)
  188. mod_path = packdir2 + TESTMOD
  189. mod_name = module_path_to_dotted_name(mod_path)
  190. pkg = __import__(mod_name)
  191. mod = sys.modules[mod_name]
  192. self.assertEquals(zi.get_source(TESTPACK), None)
  193. self.assertEquals(zi.get_source(mod_path), None)
  194. self.assertEquals(zi._get_filename(mod_path), mod.__file__)
  195. # To pass in the module name instead of the path, we must use the right importer
  196. loader = mod.__loader__
  197. self.assertEquals(loader.get_source(mod_name), None)
  198. self.assertEquals(loader._get_filename(mod_name), mod.__file__)
  199. # test prefix and archivepath members
  200. zi2 = zipimport.zipimporter(TEMP_ZIP + os.sep + TESTPACK)
  201. self.assertEquals(zi2.archive, TEMP_ZIP)
  202. self.assertEquals(zi2.prefix, TESTPACK + os.sep)
  203. finally:
  204. z.close()
  205. os.remove(TEMP_ZIP)
  206. def testZipImporterMethodsInSubDirectory(self):
  207. packdir = TESTPACK + os.sep
  208. packdir2 = packdir + TESTPACK2 + os.sep
  209. files = {packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
  210. packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
  211. z = ZipFile(TEMP_ZIP, "w")
  212. try:
  213. for name, (mtime, data) in files.items():
  214. zinfo = ZipInfo(name, time.localtime(mtime))
  215. zinfo.compress_type = self.compression
  216. z.writestr(zinfo, data)
  217. z.close()
  218. zi = zipimport.zipimporter(TEMP_ZIP + os.sep + packdir)
  219. self.assertEquals(zi.archive, TEMP_ZIP)
  220. self.assertEquals(zi.prefix, packdir)
  221. self.assertEquals(zi.is_package(TESTPACK2), True)
  222. mod = zi.load_module(TESTPACK2)
  223. self.assertEquals(zi._get_filename(TESTPACK2), mod.__file__)
  224. self.assertEquals(zi.is_package(TESTPACK2 + os.sep + '__init__'), False)
  225. self.assertEquals(zi.is_package(TESTPACK2 + os.sep + TESTMOD), False)
  226. mod_path = TESTPACK2 + os.sep + TESTMOD
  227. mod_name = module_path_to_dotted_name(mod_path)
  228. pkg = __import__(mod_name)
  229. mod = sys.modules[mod_name]
  230. self.assertEquals(zi.get_source(TESTPACK2), None)
  231. self.assertEquals(zi.get_source(mod_path), None)
  232. self.assertEquals(zi._get_filename(mod_path), mod.__file__)
  233. # To pass in the module name instead of the path, we must use the right importer
  234. loader = mod.__loader__
  235. self.assertEquals(loader.get_source(mod_name), None)
  236. self.assertEquals(loader._get_filename(mod_name), mod.__file__)
  237. finally:
  238. z.close()
  239. os.remove(TEMP_ZIP)
  240. def testGetData(self):
  241. z = ZipFile(TEMP_ZIP, "w")
  242. z.compression = self.compression
  243. try:
  244. name = "testdata.dat"
  245. data = "".join([chr(x) for x in range(256)]) * 500
  246. z.writestr(name, data)
  247. z.close()
  248. zi = zipimport.zipimporter(TEMP_ZIP)
  249. self.assertEquals(data, zi.get_data(name))
  250. self.assert_('zipimporter object' in repr(zi))
  251. finally:
  252. z.close()
  253. os.remove(TEMP_ZIP)
  254. def testImporterAttr(self):
  255. src = """if 1: # indent hack
  256. def get_file():
  257. return __file__
  258. if __loader__.get_data("some.data") != "some data":
  259. raise AssertionError, "bad data"\n"""
  260. pyc = make_pyc(compile(src, "<???>", "exec"), NOW)
  261. files = {TESTMOD + pyc_ext: (NOW, pyc),
  262. "some.data": (NOW, "some data")}
  263. self.doTest(pyc_ext, files, TESTMOD)
  264. def testImport_WithStuff(self):
  265. # try importing from a zipfile which contains additional
  266. # stuff at the beginning of the file
  267. files = {TESTMOD + ".py": (NOW, test_src)}
  268. self.doTest(".py", files, TESTMOD,
  269. stuff="Some Stuff"*31)
  270. def assertModuleSource(self, module):
  271. self.assertEqual(inspect.getsource(module), test_src)
  272. def testGetSource(self):
  273. files = {TESTMOD + ".py": (NOW, test_src)}
  274. self.doTest(".py", files, TESTMOD, call=self.assertModuleSource)
  275. def testGetCompiledSource(self):
  276. pyc = make_pyc(compile(test_src, "<???>", "exec"), NOW)
  277. files = {TESTMOD + ".py": (NOW, test_src),
  278. TESTMOD + pyc_ext: (NOW, pyc)}
  279. self.doTest(pyc_ext, files, TESTMOD, call=self.assertModuleSource)
  280. def runDoctest(self, callback):
  281. files = {TESTMOD + ".py": (NOW, test_src),
  282. "xyz.txt": (NOW, ">>> log.append(True)\n")}
  283. self.doTest(".py", files, TESTMOD, call=callback)
  284. def doDoctestFile(self, module):
  285. log = []
  286. old_master, doctest.master = doctest.master, None
  287. try:
  288. doctest.testfile(
  289. 'xyz.txt', package=module, module_relative=True,
  290. globs=locals()
  291. )
  292. finally:
  293. doctest.master = old_master
  294. self.assertEqual(log,[True])
  295. def testDoctestFile(self):
  296. self.runDoctest(self.doDoctestFile)
  297. def doDoctestSuite(self, module):
  298. log = []
  299. doctest.DocFileTest(
  300. 'xyz.txt', package=module, module_relative=True,
  301. globs=locals()
  302. ).run()
  303. self.assertEqual(log,[True])
  304. def testDoctestSuite(self):
  305. self.runDoctest(self.doDoctestSuite)
  306. def doTraceback(self, module):
  307. try:
  308. module.do_raise()
  309. except:
  310. tb = sys.exc_info()[2].tb_next
  311. f,lno,n,line = extract_tb(tb, 1)[0]
  312. self.assertEqual(line, raise_src.strip())
  313. f,lno,n,line = extract_stack(tb.tb_frame, 1)[0]
  314. self.assertEqual(line, raise_src.strip())
  315. s = StringIO.StringIO()
  316. print_tb(tb, 1, s)
  317. self.failUnless(s.getvalue().endswith(raise_src))
  318. else:
  319. raise AssertionError("This ought to be impossible")
  320. def testTraceback(self):
  321. files = {TESTMOD + ".py": (NOW, raise_src)}
  322. self.doTest(None, files, TESTMOD, call=self.doTraceback)
  323. class CompressedZipImportTestCase(UncompressedZipImportTestCase):
  324. compression = ZIP_DEFLATED
  325. class BadFileZipImportTestCase(unittest.TestCase):
  326. def assertZipFailure(self, filename):
  327. self.assertRaises(zipimport.ZipImportError,
  328. zipimport.zipimporter, filename)
  329. def testNoFile(self):
  330. self.assertZipFailure('AdfjdkFJKDFJjdklfjs')
  331. def testEmptyFilename(self):
  332. self.assertZipFailure('')
  333. def testBadArgs(self):
  334. self.assertRaises(TypeError, zipimport.zipimporter, None)
  335. self.assertRaises(TypeError, zipimport.zipimporter, TESTMOD, kwd=None)
  336. def testFilenameTooLong(self):
  337. self.assertZipFailure('A' * 33000)
  338. def testEmptyFile(self):
  339. test_support.unlink(TESTMOD)
  340. open(TESTMOD, 'w+').close()
  341. self.assertZipFailure(TESTMOD)
  342. def testFileUnreadable(self):
  343. test_support.unlink(TESTMOD)
  344. fd = os.open(TESTMOD, os.O_CREAT, 000)
  345. try:
  346. os.close(fd)
  347. self.assertZipFailure(TESTMOD)
  348. finally:
  349. # If we leave "the read-only bit" set on Windows, nothing can
  350. # delete TESTMOD, and later tests suffer bogus failures.
  351. os.chmod(TESTMOD, 0666)
  352. test_support.unlink(TESTMOD)
  353. def testNotZipFile(self):
  354. test_support.unlink(TESTMOD)
  355. fp = open(TESTMOD, 'w+')
  356. fp.write('a' * 22)
  357. fp.close()
  358. self.assertZipFailure(TESTMOD)
  359. # XXX: disabled until this works on Big-endian machines
  360. def _testBogusZipFile(self):
  361. test_support.unlink(TESTMOD)
  362. fp = open(TESTMOD, 'w+')
  363. fp.write(struct.pack('=I', 0x06054B50))
  364. fp.write('a' * 18)
  365. fp.close()
  366. z = zipimport.zipimporter(TESTMOD)
  367. try:
  368. self.assertRaises(TypeError, z.find_module, None)
  369. self.assertRaises(TypeError, z.load_module, None)
  370. self.assertRaises(TypeError, z.is_package, None)
  371. self.assertRaises(TypeError, z.get_code, None)
  372. self.assertRaises(TypeError, z.get_data, None)
  373. self.assertRaises(TypeError, z.get_source, None)
  374. error = zipimport.ZipImportError
  375. self.assertEqual(z.find_module('abc'), None)
  376. self.assertRaises(error, z.load_module, 'abc')
  377. self.assertRaises(error, z.get_code, 'abc')
  378. self.assertRaises(IOError, z.get_data, 'abc')
  379. self.assertRaises(error, z.get_source, 'abc')
  380. self.assertRaises(error, z.is_package, 'abc')
  381. finally:
  382. zipimport._zip_directory_cache.clear()
  383. def cleanup():
  384. # this is necessary if test is run repeated (like when finding leaks)
  385. global test_imported
  386. if test_imported:
  387. zipimport._zip_directory_cache.clear()
  388. if hasattr(UncompressedZipImportTestCase, 'testAFakeZlib'):
  389. delattr(UncompressedZipImportTestCase, 'testAFakeZlib')
  390. if hasattr(CompressedZipImportTestCase, 'testAFakeZlib'):
  391. delattr(CompressedZipImportTestCase, 'testAFakeZlib')
  392. test_imported = True
  393. def test_main():
  394. cleanup()
  395. try:
  396. test_support.run_unittest(
  397. UncompressedZipImportTestCase,
  398. CompressedZipImportTestCase,
  399. BadFileZipImportTestCase,
  400. )
  401. finally:
  402. test_support.unlink(TESTMOD)
  403. if __name__ == "__main__":
  404. test_main()