loaders WIP

ansible_mitogen/loaders.py

@@ -99,5 +99,5 @@ def assert_supported_release():
 from ansible.plugins.loader import strategy_loader
 
 # These are original, unwrapped implementations
-action_loader__get = action_loader.get
-connection_loader__get = connection_loader.get_with_context
+action_loader__get_with_context = action_loader.get_with_context
+connection_loader__get_with_context = connection_loader.get_with_context

ansible_mitogen/plugins/connection/mitogen_kubectl.py

@@ -46,12 +46,6 @@
 import ansible_mitogen.loaders
 
 
-_get_result = ansible_mitogen.loaders.connection_loader__get(
-    'kubectl',
-    class_only=True,
-)
-
-
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'kubectl'
 
@@ -61,17 +55,22 @@ class Connection(ansible_mitogen.connection.Connection):
     )
 
     def __init__(self, *args, **kwargs):
-        if not _get_result:
+        result = ansible_mitogen.loaders.connection_loader__get_with_context(
+            'kubectl', class_only=True,
+        )
+        klass = result.object
+        if not klass:
             raise ansible.errors.AnsibleConnectionFailure(self.not_supported_msg)
+
+        # FIXME Better name
+        # TODO Generalize?
+        self._klass = klass
         super(Connection, self).__init__(*args, **kwargs)
 
     def get_extra_args(self):
-        try:
-            # Ansible < 2.10, _get_result is the connection class
-            connection_options = _get_result.connection_options
-        except AttributeError:
-            # Ansible >= 2.10, _get_result is a get_with_context_result
-            connection_options = _get_result.object.connection_options
+        # TODO Can this overridden method be replaced by similar to used in mitogen_ssh?
+        # TODO Can self._klass.connection_options be a basis for something better than used in mitogen_ssh?
+        connection_options = self._klass.connection_options
         parameters = []
         for key in connection_options:
             task_var_name = 'ansible_%s' % key

ansible_mitogen/plugins/connection/mitogen_ssh.py

@@ -33,6 +33,7 @@
 import sys
 
 from ansible.plugins.connection.ssh import (
+    Connection as _ansible_ssh_Connection,
     DOCUMENTATION as _ansible_ssh_DOCUMENTATION,
 )
 
@@ -55,18 +56,13 @@
     del base_dir
 
 import ansible_mitogen.connection
-import ansible_mitogen.loaders
 
 
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'ssh'
-    vanilla_class = ansible_mitogen.loaders.connection_loader__get(
-        'ssh',
-        class_only=True,
-    )
 
     @staticmethod
     def _create_control_path(*args, **kwargs):
         """Forward _create_control_path() to the implementation in ssh.py."""
         # https://github.com/dw/mitogen/issues/342
-        return Connection.vanilla_class._create_control_path(*args, **kwargs)
+        return _ansible_ssh_Connection._create_control_path(*args, **kwargs)

ansible_mitogen/strategy.py

@@ -76,7 +76,7 @@ def patch_add_local(self, **kwargs):
 _patch_awx_callback()
 
 
-def wrap_action_loader__get(name, *args, **kwargs):
+def wrap_action_loader__get_with_context(name, class_only=False, collection_list=None, *args, **kwargs):
     """
     While the mitogen strategy is active, trap action_loader.get() calls,
     augmenting any fetched class with ActionModuleMixin, which replaces various
@@ -86,18 +86,23 @@ def wrap_action_loader__get(name, *args, **kwargs):
     This is used instead of static subclassing as it generalizes to third party
     action plugins outside the Ansible tree.
     """
-    get_kwargs = {'class_only': True}
     if name in ('fetch',):
         name = 'mitogen_' + name
-    get_kwargs['collection_list'] = kwargs.pop('collection_list', None)
 
-    klass = ansible_mitogen.loaders.action_loader__get(name, **get_kwargs)
-    if klass:
-        bases = (ansible_mitogen.mixins.ActionModuleMixin, klass)
-        adorned_klass = type(str(name), bases, {})
-        if kwargs.get('class_only'):
-            return adorned_klass
-        return adorned_klass(*args, **kwargs)
+    result = ansible_mitogen.loaders.action_loader__get_with_context(
+        name, class_only=True, collection_list=collection_list,
+    )
+
+    if not result.object:
+        return result
+
+    adorned_bases = (ansible_mitogen.mixins.ActionModuleMixin, result.object)
+    adorned_class = type(str(name), adorned_bases, {})
+
+    if class_only:
+        return result._replace(object=adorned_class)
+
+    return result._replace(object=adorned_class(*args, **kwargs))
 
 
 REDIRECTED_CONNECTION_PLUGINS = (
@@ -115,15 +120,17 @@ def wrap_action_loader__get(name, *args, **kwargs):
 )
 
 
-def wrap_connection_loader__get(name, *args, **kwargs):
+def wrap_connection_loader__get_with_context(name, *args, **kwargs):
     """
     While a Mitogen strategy is active, rewrite connection_loader.get() calls
     for some transports into requests for a compatible Mitogen transport.
     """
     if name in REDIRECTED_CONNECTION_PLUGINS:
         name = 'mitogen_' + name
 
-    return ansible_mitogen.loaders.connection_loader__get(name, *args, **kwargs)
+    return ansible_mitogen.loaders.connection_loader__get_with_context(
+        name, *args, **kwargs
+    )
 
 
 def wrap_worker__run(self):
@@ -173,8 +180,8 @@ def _install_wrappers(self):
         Install our PluginLoader monkey patches and update global variables
         with references to the real functions.
         """
-        ansible_mitogen.loaders.action_loader.get = wrap_action_loader__get
-        ansible_mitogen.loaders.connection_loader.get_with_context = wrap_connection_loader__get
+        ansible_mitogen.loaders.action_loader.get_with_context = wrap_action_loader__get_with_context
+        ansible_mitogen.loaders.connection_loader.get_with_context = wrap_connection_loader__get_with_context
 
         global worker__run
         worker__run = ansible.executor.process.worker.WorkerProcess.run
@@ -184,12 +191,8 @@ def _remove_wrappers(self):
         """
         Uninstall the PluginLoader monkey patches.
         """
-        ansible_mitogen.loaders.action_loader.get = (
-            ansible_mitogen.loaders.action_loader__get
-        )
-        ansible_mitogen.loaders.connection_loader.get_with_context = (
-            ansible_mitogen.loaders.connection_loader__get
-        )
+        ansible_mitogen.loaders.action_loader.get_with_context = ansible_mitogen.loaders.action_loader__get_with_context
+        ansible_mitogen.loaders.connection_loader.get_with_context = ansible_mitogen.loaders.connection_loader__get_with_context
         ansible.executor.process.worker.WorkerProcess.run = worker__run
 
     def install(self):