home *** CD-ROM | disk | FTP | other *** search
/ H4CK3R 14 / hacker14.iso / programacao / pythonwin / python.exe / TEST_ZIPIMPORT.PY < prev    next >
Encoding:
Python Source  |  2003-07-22  |  7.3 KB  |  216 lines
  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.         if "zlib" in sys.modules:
  99.             del sys.modules["zlib"]
  100.         files = {"zlib.py": (NOW, test_src)}
  101.         try:
  102.             self.doTest(".py", files, "zlib")
  103.         except ImportError:
  104.             if self.compression != ZIP_DEFLATED:
  105.                 self.fail("expected test to not raise ImportError")
  106.         else:
  107.             if self.compression != ZIP_STORED:
  108.                 self.fail("expected test to raise ImportError")
  109.  
  110.     def testPy(self):
  111.         files = {TESTMOD + ".py": (NOW, test_src)}
  112.         self.doTest(".py", files, TESTMOD)
  113.  
  114.     def testPyc(self):
  115.         files = {TESTMOD + pyc_ext: (NOW, test_pyc)}
  116.         self.doTest(pyc_ext, files, TESTMOD)
  117.  
  118.     def testBoth(self):
  119.         files = {TESTMOD + ".py": (NOW, test_src),
  120.                  TESTMOD + pyc_ext: (NOW, test_pyc)}
  121.         self.doTest(pyc_ext, files, TESTMOD)
  122.  
  123.     def testEmptyPy(self):
  124.         files = {TESTMOD + ".py": (NOW, "")}
  125.         self.doTest(None, files, TESTMOD)
  126.  
  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.  
  136.     def testBadMagic2(self):
  137.         # make pyc magic word invalid, causing an ImportError
  138.         m0 = ord(test_pyc[0])
  139.         m0 ^= 0x04  # flip an arbitrary bit
  140.         badmagic_pyc = chr(m0) + test_pyc[1:]
  141.         files = {TESTMOD + pyc_ext: (NOW, badmagic_pyc)}
  142.         try:
  143.             self.doTest(".py", files, TESTMOD)
  144.         except ImportError:
  145.             pass
  146.         else:
  147.             self.fail("expected ImportError; import from bad pyc")
  148.  
  149.     def testBadMTime(self):
  150.         t3 = ord(test_pyc[7])
  151.         t3 ^= 0x02  # flip the second bit -- not the first as that one
  152.                     # isn't stored in the .py's mtime in the zip archive.
  153.         badtime_pyc = test_pyc[:7] + chr(t3) + test_pyc[8:]
  154.         files = {TESTMOD + ".py": (NOW, test_src),
  155.                  TESTMOD + pyc_ext: (NOW, badtime_pyc)}
  156.         self.doTest(".py", files, TESTMOD)
  157.  
  158.     def testPackage(self):
  159.         packdir = TESTPACK + os.sep
  160.         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
  161.                  packdir + TESTMOD + pyc_ext: (NOW, test_pyc)}
  162.         self.doTest(pyc_ext, files, TESTPACK, TESTMOD)
  163.  
  164.     def testDeepPackage(self):
  165.         packdir = TESTPACK + os.sep
  166.         packdir2 = packdir + TESTPACK2 + os.sep
  167.         files = {packdir + "__init__" + pyc_ext: (NOW, test_pyc),
  168.                  packdir2 + "__init__" + pyc_ext: (NOW, test_pyc),
  169.                  packdir2 + TESTMOD + pyc_ext: (NOW, test_pyc)}
  170.         self.doTest(pyc_ext, files, TESTPACK, TESTPACK2, TESTMOD)
  171.  
  172.     def testGetData(self):
  173.         z = ZipFile(TEMP_ZIP, "w")
  174.         z.compression = self.compression
  175.         try:
  176.             name = "testdata.dat"
  177.             data = "".join([chr(x) for x in range(256)]) * 500
  178.             z.writestr(name, data)
  179.             z.close()
  180.             zi = zipimport.zipimporter(TEMP_ZIP)
  181.             self.assertEquals(data, zi.get_data(name))
  182.         finally:
  183.             z.close()
  184.             os.remove(TEMP_ZIP)
  185.  
  186.     def testImporterAttr(self):
  187.         src = """if 1:  # indent hack
  188.         def get_file():
  189.             return __file__
  190.         if __loader__.get_data("some.data") != "some data":
  191.             raise AssertionError, "bad data"\n"""
  192.         pyc = make_pyc(compile(src, "<???>", "exec"), NOW)
  193.         files = {TESTMOD + pyc_ext: (NOW, pyc),
  194.                  "some.data": (NOW, "some data")}
  195.         self.doTest(pyc_ext, files, TESTMOD)
  196.  
  197.     def testImport_WithStuff(self):
  198.         # try importing from a zipfile which contains additional
  199.         # stuff at the beginning of the file
  200.         files = {TESTMOD + ".py": (NOW, test_src)}
  201.         self.doTest(".py", files, TESTMOD,
  202.                     stuff="Some Stuff"*31)
  203.  
  204. class CompressedZipImportTestCase(UncompressedZipImportTestCase):
  205.     compression = ZIP_DEFLATED
  206.  
  207.  
  208. def test_main():
  209.     test_support.run_unittest(
  210.         UncompressedZipImportTestCase,
  211.         CompressedZipImportTestCase
  212.     )
  213.  
  214. if __name__ == "__main__":
  215.     test_main()
  216.