1 from test.test_support import run_unittest

     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 = '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 class PkgutilPEP302Tests(unittest.TestCase):

    82 

    83     class MyTestLoader(object):

    84         def load_module(self, fullname):

    85             # Create an empty module

    86             mod = sys.modules.setdefault(fullname, imp.new_module(fullname))

    87             mod.__file__ = "<%s>" % self.__class__.__name__

    88             mod.__loader__ = self

    89             # Make it a package

    90             mod.__path__ = []

    91             # Count how many times the module is reloaded

    92             mod.__dict__['loads'] = mod.__dict__.get('loads',0) + 1

    93             return mod

    94 

    95         def get_data(self, path):

    96             return "Hello, world!"

    97 

    98     class MyTestImporter(object):

    99         def find_module(self, fullname, path=None):

   100             return PkgutilPEP302Tests.MyTestLoader()

   101 

   102     def setUp(self):

   103         sys.meta_path.insert(0, self.MyTestImporter())

   104 

   105     def tearDown(self):

   106         del sys.meta_path[0]

   107 

   108     def test_getdata_pep302(self):

   109         # Use a dummy importer/loader

   110         self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")

   111         del sys.modules['foo']

   112 

   113     def test_alreadyloaded(self):

   114         # Ensure that get_data works without reloading - the "loads" module

   115         # variable in the example loader should count how many times a reload

   116         # occurs.

   117         import foo

   118         self.assertEqual(foo.loads, 1)

   119         self.assertEqual(pkgutil.get_data('foo', 'dummy'), "Hello, world!")

   120         self.assertEqual(foo.loads, 1)

   121         del sys.modules['foo']

   122 

   123 def test_main():

   124     run_unittest(PkgutilTests, PkgutilPEP302Tests)

   125     # this is necessary if test is run repeated (like when finding leaks)

   126     import zipimport

   127     zipimport._zip_directory_cache.clear()

   128 

   129 if __name__ == '__main__':

   130     test_main()