From 4000dd762827d734528fce307ea890b483f8c5e6 Mon Sep 17 00:00:00 2001
From: David Shrewsbury <shrewsbury.dave@gmail.com>
Date: Fri, 21 Sep 2012 15:06:59 -0400
Subject: [PATCH] Add code to dynamically import a class.

---
 .gitignore                   |  1 +
 libra/worker/drivers/base.py |  6 ++++++
 libra/worker/utils.py        | 29 +++++++++++++++++++++++++++++
 libra/worker/worker.py       | 19 +++++++++++++++----
 4 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 libra/worker/utils.py

diff --git a/.gitignore b/.gitignore
index b9c73dba..b0f9693f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@
 dist
 build
 .cache
+*.swp
diff --git a/libra/worker/drivers/base.py b/libra/worker/drivers/base.py
index 6bc71785..d1a75570 100644
--- a/libra/worker/drivers/base.py
+++ b/libra/worker/drivers/base.py
@@ -12,6 +12,12 @@
 # License for the specific language governing permissions and limitations
 
 
+# Mapping of --driver options to a class
+known_drivers = {
+    'haproxy': 'libra.worker.drivers.haproxy.driver.HAProxyDriver'
+}
+
+
 class LoadBalancerDriver(object):
     """
     Load balancer device driver base class.
diff --git a/libra/worker/utils.py b/libra/worker/utils.py
new file mode 100644
index 00000000..1c0f55a1
--- /dev/null
+++ b/libra/worker/utils.py
@@ -0,0 +1,29 @@
+# Copyright 2012 Hewlett-Packard Development Company, L.P.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import importlib
+import sys
+import traceback
+
+
+def import_class(import_str):
+    """ Returns a class from a string including module and class """
+    mod_str, _sep, class_str = import_str.rpartition('.')
+    try:
+        mod = importlib.import_module(mod_str, 'libra.worker.drivers')
+        return getattr(mod, class_str)
+    except (ValueError, AttributeError):
+        raise ImportError('Class %s cannot be found (%s)' %
+                          (class_str,
+                           traceback.format_exception(*sys.exc_info())))
diff --git a/libra/worker/worker.py b/libra/worker/worker.py
index 4efdcc91..d9b14f2a 100644
--- a/libra/worker/worker.py
+++ b/libra/worker/worker.py
@@ -22,6 +22,8 @@ from time import sleep
 from libra.common.json_gearman import JSONGearmanWorker
 from libra.common.faults import BadRequest
 from libra.common.options import Options, setup_logging
+from libra.worker.drivers.base import known_drivers
+from libra.worker.utils import import_class
 
 
 def lbaas_task(worker, job):
@@ -81,6 +83,7 @@ def lbaas_task(worker, job):
 
 
 class CustomJSONGearmanWorker(JSONGearmanWorker):
+    """ Custom class we will use to pass arguments to the Gearman task. """
     logger = None
     driver = None
 
@@ -133,14 +136,22 @@ def main():
         default=60, help='seconds to sleep between job server reconnects'
     )
     options.parser.add_argument(
-        '--driver', dest='driver', default='haproxy.driver.HAProxyDriver',
-        help='Class name of device driver to use'
+        '--driver', dest='driver',
+        choices=known_drivers.keys(), default='haproxy',
+        help='Type of device to use'
     )
     args = options.run()
 
     logger = setup_logging('libra_worker', args)
-    from libra.worker.drivers.haproxy.driver import HAProxyDriver
-    driver = HAProxyDriver()
+
+    # Import the device driver we are going to use. This will be sent
+    # along to the Gearman task that will use it to communicate with
+    # the device.
+
+    logger.debug("Using driver %s=%s" % (args.driver,
+                                         known_drivers[args.driver]))
+    driver_class = import_class(known_drivers[args.driver])
+    driver = driver_class()
 
     server = Server(['localhost:4730'], args.reconnect_sleep)
     server.logger = logger