Merge pull request #20481 from jacksontj/loader

Add submodule support to LazyLoader

salt/loader.py

@@ -710,6 +710,7 @@ class LazyLoader(salt.utils.lazy.LazyDict):
 
         self.whitelist = whitelist
         self.virtual_enable = virtual_enable
+        self.initial_load = True
 
         # names of modules that we don't have (errors, __virtual__, etc.)
         self.missing_modules = []
@@ -745,6 +746,8 @@ class LazyLoader(salt.utils.lazy.LazyDict):
             except ImportError:
                 log.info('Cython is enabled in the options but not present '
                     'in the system path. Skipping Cython modules.')
+        # allow for module dirs
+        self.suffix_map[''] = ('', '', imp.PKG_DIRECTORY)
 
         # create mapping of filename (without suffix) to (path, suffix)
         self.file_mapping = {}
@@ -768,8 +771,21 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                         )
                         continue
                     fpath = os.path.join(mod_dir, filename)
+                    # if its a directory, lets allow us to load that
+                    if ext == '':
+                        # is there something __init__?
+                        subfiles = os.listdir(fpath)
+                        sub_path = None
+                        for suffix in suffix_order:
+                            init_file = '__init__{0}'.format(suffix)
+                            if init_file in subfiles:
+                                sub_path = os.path.join(fpath, init_file)
+                                break
+                        if sub_path is not None:
+                            self.file_mapping[f_noext] = (fpath, ext)
+
                     # if we don't have it, we want it
-                    if f_noext not in self.file_mapping:
+                    elif f_noext not in self.file_mapping:
                         self.file_mapping[f_noext] = (fpath, ext)
                     # if we do, we want it if we have a higher precidence ext
                     else:
@@ -790,6 +806,7 @@ class LazyLoader(salt.utils.lazy.LazyDict):
         # we obviously want a re-do
         if hasattr(self, 'opts'):
             self.refresh_file_mapping()
+        self.initial_load = False
 
     def __prep_mod_opts(self, opts):
         '''
@@ -853,6 +870,19 @@ class LazyLoader(salt.utils.lazy.LazyDict):
             if mod_name not in k:
                 yield k
 
+    def _reload_submodules(self, mod):
+        submodules = (
+            getattr(mod, sname) for sname in dir(mod) if
+            isinstance(getattr(mod, sname), mod.__class__)
+        )
+
+        # reload only custom "sub"modules
+        for submodule in submodules:
+            # it is a submodule if the name is in a namespace under mod
+            if submodule.__name__.startswith(mod.__name__ + '.'):
+                reload(submodule)
+                self._reload_submodules(submodule)
+
     def _load_module(self, name):
         mod = None
         fpath, suffix = self.file_mapping[name]
@@ -862,14 +892,27 @@ class LazyLoader(salt.utils.lazy.LazyDict):
                 mod = self.pyximport.load_module(name, fpath, tempfile.gettempdir())
             else:
                 desc = self.suffix_map[suffix]
-                with open(fpath, desc[1]) as fn_:
+                # if it is a directory, we dont open a file
+                if suffix == '':
                     mod = imp.load_module(
                         '{0}.{1}.{2}.{3}'.format(
                             self.loaded_base_name,
                             self.mod_type_check(fpath),
                             self.tag,
                             name
-                        ), fn_, fpath, desc)
+                        ), None, fpath, desc)
+                    # reload all submodules if necessary
+                    if not self.initial_load:
+                        self._reload_submodules(mod)
+                else:
+                    with open(fpath, desc[1]) as fn_:
+                        mod = imp.load_module(
+                            '{0}.{1}.{2}.{3}'.format(
+                                self.loaded_base_name,
+                                self.mod_type_check(fpath),
+                                self.tag,
+                                name
+                            ), fn_, fpath, desc)
 
         except IOError:
             raise

tests/integration/loader/loader.py

@@ -12,6 +12,7 @@ import inspect
 import tempfile
 import shutil
 import os
+import collections
 
 # Import Salt Testing libs
 from salttesting import TestCase
@@ -279,3 +280,238 @@ class LazyLoaderReloadingTest(TestCase):
         self.assertEqual(self.loader[self.module_key](), self.count)
         self.loader.clear()
         self.assertNotIn(self.module_key, self.loader)
+
+submodule_template = '''
+import lib
+
+def test():
+    return ({count}, lib.test())
+'''
+
+submodule_lib_template = '''
+def test():
+    return {count}
+'''
+
+
+class LazyLoaderSubmodReloadingTest(TestCase):
+    '''
+    Test the loader of salt with changing modules
+    '''
+    module_name = 'loadertestsubmod'
+    module_key = 'loadertestsubmod.test'
+
+    def setUp(self):
+        self.opts = _config = minion_config(None)
+        self.tmp_dir = tempfile.mkdtemp(dir=integration.TMP)
+        os.makedirs(self.module_dir)
+
+        self.count = 0
+        self.lib_count = 0
+
+        dirs = _module_dirs(self.opts, 'modules', 'module')
+        dirs.append(self.tmp_dir)
+        self.loader = LazyLoader(dirs,
+                                 self.opts,
+                                 tag='module')
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def update_module(self):
+        self.count += 1
+        with open(self.module_path, 'wb') as fh:
+            fh.write(submodule_template.format(count=self.count))
+            fh.flush()
+            os.fsync(fh.fileno())  # flush to disk
+
+        # pyc files don't like it when we change the original quickly
+        # since the header bytes only contain the timestamp (granularity of seconds)
+        # TODO: don't write them? Is *much* slower on re-load (~3x)
+        # https://docs.python.org/2/library/sys.html#sys.dont_write_bytecode
+        try:
+            os.unlink(self.module_path + 'c')
+        except OSError:
+            pass
+
+    def rm_module(self):
+        os.unlink(self.module_path)
+        os.unlink(self.module_path + 'c')
+
+    def update_lib(self):
+        self.lib_count += 1
+        with open(self.lib_path, 'wb') as fh:
+            fh.write(submodule_lib_template.format(count=self.lib_count))
+            fh.flush()
+            os.fsync(fh.fileno())  # flush to disk
+
+        # pyc files don't like it when we change the original quickly
+        # since the header bytes only contain the timestamp (granularity of seconds)
+        # TODO: don't write them? Is *much* slower on re-load (~3x)
+        # https://docs.python.org/2/library/sys.html#sys.dont_write_bytecode
+        try:
+            os.unlink(self.lib_path + 'c')
+        except OSError:
+            pass
+
+    def rm_lib(self):
+        os.unlink(self.lib_path)
+        os.unlink(self.lib_path + 'c')
+
+    @property
+    def module_dir(self):
+        return os.path.join(self.tmp_dir, self.module_name)
+
+    @property
+    def module_path(self):
+        return os.path.join(self.module_dir, '__init__.py')
+
+    @property
+    def lib_path(self):
+        return os.path.join(self.module_dir, 'lib.py')
+
+    def test_basic(self):
+        # ensure it doesn't exist
+        self.assertNotIn(self.module_key, self.loader)
+
+        self.update_module()
+        self.update_lib()
+        self.loader.clear()
+        self.assertIn(self.module_key, self.loader)
+
+    def test_reload(self):
+        # ensure it doesn't exist
+        self.assertNotIn(self.module_key, self.loader)
+
+        # update both the module and the lib
+        for x in range(1, 3):
+            self.update_module()
+            self.update_lib()
+            self.loader.clear()
+            self.assertEqual(self.loader[self.module_key](), (self.count, self.lib_count))
+
+        # update just the module
+        for x in range(1, 3):
+            self.update_module()
+            self.loader.clear()
+            self.assertEqual(self.loader[self.module_key](), (self.count, self.lib_count))
+
+        # update just the lib
+        for x in range(1, 3):
+            self.update_lib()
+            self.loader.clear()
+            self.assertEqual(self.loader[self.module_key](), (self.count, self.lib_count))
+
+        self.rm_module()
+        # make sure that even if we remove the module, its still loaded until a clear
+        self.assertEqual(self.loader[self.module_key](), (self.count, self.lib_count))
+        self.loader.clear()
+        self.assertNotIn(self.module_key, self.loader)
+
+    def test_reload_missing_lib(self):
+        # ensure it doesn't exist
+        self.assertNotIn(self.module_key, self.loader)
+
+        # update both the module and the lib
+        self.update_module()
+        self.update_lib()
+        self.loader.clear()
+        self.assertEqual(self.loader[self.module_key](), (self.count, self.lib_count))
+
+        # remove the lib, this means we should fail to load the module next time
+        self.rm_lib()
+        self.loader.clear()
+        self.assertNotIn(self.module_key, self.loader)
+
+
+deep_init_base = '''
+import top_lib
+import top_lib.mid_lib
+import top_lib.mid_lib.bot_lib
+
+def top():
+    return top_lib.test()
+
+def mid():
+    return top_lib.mid_lib.test()
+
+def bot():
+    return top_lib.mid_lib.bot_lib.test()
+'''
+
+
+class LazyLoaderDeepSubmodReloadingTest(TestCase):
+    module_name = 'loadertestsubmoddeep'
+    libs = ('top_lib', 'mid_lib', 'bot_lib')
+
+    def setUp(self):
+        self.opts = _config = minion_config(None)
+        self.tmp_dir = tempfile.mkdtemp(dir=integration.TMP)
+        os.makedirs(self.module_dir)
+
+        self.lib_count = collections.defaultdict(int)  # mapping of path -> count
+
+        # bootstrap libs
+        with open(os.path.join(self.module_dir, '__init__.py'), 'w') as fh:
+            fh.write(deep_init_base)
+            fh.flush()
+            os.fsync(fh.fileno())  # flush to disk
+
+        self.lib_paths = {}
+        dir_path = self.module_dir
+        for lib_name in self.libs:
+            dir_path = os.path.join(dir_path, lib_name)
+            self.lib_paths[lib_name] = dir_path
+            os.makedirs(dir_path)
+            self.update_lib(lib_name)
+
+        dirs = _module_dirs(self.opts, 'modules', 'module')
+        dirs.append(self.tmp_dir)
+        self.loader = LazyLoader(dirs,
+                                 self.opts,
+                                 tag='module')
+
+    @property
+    def module_dir(self):
+        return os.path.join(self.tmp_dir, self.module_name)
+
+    def update_lib(self, lib_name):
+        path = os.path.join(self.lib_paths[lib_name], '__init__.py')
+        self.lib_count[lib_name] += 1
+        with open(path, 'wb') as fh:
+            fh.write(submodule_lib_template.format(count=self.lib_count[lib_name]))
+            fh.flush()
+            os.fsync(fh.fileno())  # flush to disk
+
+        # pyc files don't like it when we change the original quickly
+        # since the header bytes only contain the timestamp (granularity of seconds)
+        # TODO: don't write them? Is *much* slower on re-load (~3x)
+        # https://docs.python.org/2/library/sys.html#sys.dont_write_bytecode
+        try:
+            os.unlink(path + 'c')
+        except OSError:
+            pass
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+
+    def test_basic(self):
+        self.assertIn('{0}.top'.format(self.module_name), self.loader)
+
+    def _verify_libs(self):
+        for lib in self.libs:
+            self.assertEqual(self.loader['{0}.{1}'.format(self.module_name, lib.replace('_lib', ''))](),
+                             self.lib_count[lib])
+
+    def test_reload(self):
+        '''
+        Make sure that we can reload all libraries of arbitrary depth
+        '''
+        self._verify_libs()
+
+        # update them all
+        for lib in self.libs:
+            for x in xrange(5):
+                self.update_lib(lib)
+                self.loader.clear()
+                self._verify_libs()