Merge pull request #320 from napalm-automation/merge-master

Porting the commits from #319 into the master branch
napalm-automation · Nov 23, 2017 · e462990 · e462990
2 parents 1150ca4 + 16d20f9
commit e462990
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 71 deletions.
diff --git a/napalm_base/__init__.py b/napalm_base/__init__.py
@@ -20,8 +20,6 @@
 
 # Python std lib
 import sys
-import inspect
-import importlib
 import pkg_resources
 
 # Verify Python Version that is running
@@ -32,83 +30,106 @@
 except AttributeError:
     raise RuntimeError('NAPALM requires Python 2.7 or Python3')
 
-# NAPALM base
-from napalm_base.base import NetworkDriver
-from napalm_base.exceptions import ModuleImportError
-from napalm_base.mock import MockDriver
-from napalm_base.utils import py23_compat
+# Try to import napalm
+try:
+    import napalm
+    HAS_NAPALM = True
+    try:
+        NAPALM_MAJOR = int(napalm.__version__.split('.')[0])
+    except AttributeError:
+        NAPALM_MAJOR = 0
+except ImportError:
+    HAS_NAPALM = False
 
 try:
     __version__ = pkg_resources.get_distribution('napalm-base').version
 except pkg_resources.DistributionNotFound:
     __version__ = "Not installed"
 
 
-__all__ = [
-    'get_network_driver',  # export the function
-    'NetworkDriver'  # also export the base class
-]
-
+if HAS_NAPALM and NAPALM_MAJOR >= 2:
+    # If napalm >= 2.0.0 is installed, then import get_network_driver
+    from napalm import get_network_driver
+    from napalm.base import NetworkDriver
+else:
+    # Import std lib
+    import inspect
+    import importlib
+
+    # Import local modules
+    from napalm_base.exceptions import ModuleImportError
+    from napalm_base.mock import MockDriver
+    from napalm_base.utils import py23_compat
+    from napalm_base.base import NetworkDriver
+
+    def get_network_driver(module_name, prepend=True):
+        """
+        Searches for a class derived form the base NAPALM class NetworkDriver in a specific library.
+        The library name must repect the following pattern: napalm_[DEVICE_OS].
+        NAPALM community supports a list of devices and provides the corresponding libraries; for
+        full reference please refer to the `Supported Network Operation Systems`_ paragraph on
+        `Read the Docs`_.
+
+        .. _`Supported Network Operation Systems`: \
+        http://napalm.readthedocs.io/en/latest/#supported-network-operating-systems
+        .. _`Read the Docs`: \
+        http://napalm.readthedocs.io/
+
+        module_name
+            The name of the device operating system or the name of the library.
+
+        :return: The first class derived from NetworkDriver, found in the library.
+
+        :raise ModuleImportError: When the library is not installed or a derived class from \
+        NetworkDriver was not found.
+
+        Example:
+
+        .. code-block:: python
+
+            >>> get_network_driver('junos')
+            <class 'napalm_junos.junos.JunOSDriver'>
+            >>> get_network_driver('IOS-XR')
+            <class 'napalm_iosxr.iosxr.IOSXRDriver'>
+            >>> get_network_driver('napalm_eos')
+            <class 'napalm_eos.eos.EOSDriver'>
+            >>> get_network_driver('wrong')
+            napalm_base.exceptions.ModuleImportError: Cannot import "napalm_wrong". Is the library \
+            installed?
+        """
+        if module_name == "mock":
+            return MockDriver
+
+        if not (isinstance(module_name, py23_compat.string_types) and len(module_name) > 0):
+            raise ModuleImportError('Please provide a valid driver name.')
+
+        try:
+            # Only lowercase allowed
+            module_name = module_name.lower()
+            # Try to not raise error when users requests IOS-XR for e.g.
+            module_install_name = module_name.replace('-', '')
+            # Can also request using napalm_[SOMETHING]
+            if 'napalm_' not in module_install_name and prepend is True:
+                module_install_name = 'napalm_{name}'.format(name=module_install_name)
+            module = importlib.import_module(module_install_name)
+        except ImportError:
+            raise ModuleImportError(
+                    'Cannot import "{install_name}". Is the library installed?'.format(
+                        install_name=module_install_name
+                    )
+                )
 
-def get_network_driver(module_name, prepend=True):
-    """
-    Searches for a class derived form the base NAPALM class NetworkDriver in a specific library.
-    The library name must repect the following pattern: napalm_[DEVICE_OS].
-    NAPALM community supports a list of devices and provides the corresponding libraries; for
-    full reference please refer to the `Supported Network Operation Systems`_ paragraph on
-    `Read the Docs`_.
-
-    .. _`Supported Network Operation Systems`: \
-    http://napalm.readthedocs.io/en/latest/#supported-network-operating-systems
-    .. _`Read the Docs`: \
-    http://napalm.readthedocs.io/
-
-    :param module_name:         the name of the device operating system or the name of the library.
-    :return:                    the first class derived from NetworkDriver, found in the library.
-    :raise ModuleImportError:   when the library is not installed or a derived class from \
-    NetworkDriver was not found.
-
-    Example::
-
-    .. code-block:: python
-
-        >>> get_network_driver('junos')
-        <class 'napalm_junos.junos.JunOSDriver'>
-        >>> get_network_driver('IOS-XR')
-        <class 'napalm_iosxr.iosxr.IOSXRDriver'>
-        >>> get_network_driver('napalm_eos')
-        <class 'napalm_eos.eos.EOSDriver'>
-        >>> get_network_driver('wrong')
-        napalm_base.exceptions.ModuleImportError: Cannot import "napalm_wrong". Is the library \
-        installed?
-    """
-    if module_name == "mock":
-        return MockDriver
-
-    if not (isinstance(module_name, py23_compat.string_types) and len(module_name) > 0):
-        raise ModuleImportError('Please provide a valid driver name.')
+        for name, obj in inspect.getmembers(module):
+            if inspect.isclass(obj) and issubclass(obj, NetworkDriver):
+                return obj
 
-    try:
-        # Only lowercase allowed
-        module_name = module_name.lower()
-        # Try to not raise error when users requests IOS-XR for e.g.
-        module_install_name = module_name.replace('-', '')
-        # Can also request using napalm_[SOMETHING]
-        if 'napalm_' not in module_install_name and prepend is True:
-            module_install_name = 'napalm_{name}'.format(name=module_install_name)
-        module = importlib.import_module(module_install_name)
-    except ImportError:
+        # looks like you don't have any Driver class in your module...
         raise ModuleImportError(
-                'Cannot import "{install_name}". Is the library installed?'.format(
-                    install_name=module_install_name
-                )
-            )
+            'No class inheriting "napalm_base.base.NetworkDriver" found in "{install_name}".'
+            .format(install_name=module_install_name))
 
-    for name, obj in inspect.getmembers(module):
-        if inspect.isclass(obj) and issubclass(obj, NetworkDriver):
-            return obj
 
-    # looks like you don't have any Driver class in your module...
-    raise ModuleImportError(
-        'No class inheriting "napalm_base.base.NetworkDriver" found in "{install_name}".'
-        .format(install_name=module_install_name))
+__all__ = [
+    'get_network_driver',  # export the function
+    'NetworkDriver'  # also export the base class
+]
diff --git a/setup.py b/setup.py
@@ -12,7 +12,7 @@
 
 setup(
     name="napalm-base",
-    version='0.25.0',
+    version='1.0.0',
     packages=find_packages(),
     author="David Barroso, Kirk Byers, Mircea Ulinic",
     author_email="[email protected], [email protected], [email protected]",