From 16d20f9127cc864c95f585fd7aa73584e2dd35c7 Mon Sep 17 00:00:00 2001
From: Mircea Ulinic <mirucha@cloudflare.com>
Date: Wed, 15 Nov 2017 16:24:38 +0000
Subject: [PATCH] Import NetworkDriver from napalm.base when napalm>=2.0.0 is
 installed

Otherwise, couple together the local napalm_base
get_network_driver and NetworkDriver.
---
 napalm_base/__init__.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/napalm_base/__init__.py b/napalm_base/__init__.py
index e04cbaae..6a82fb41 100644
--- a/napalm_base/__init__.py
+++ b/napalm_base/__init__.py
@@ -22,9 +22,6 @@
 import sys
 import pkg_resources
 
-# Import local modules
-from napalm_base.base import NetworkDriver
-
 # Verify Python Version that is running
 try:
     if not(sys.version_info.major == 2 and sys.version_info.minor == 7) and \
@@ -53,6 +50,7 @@
 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
@@ -62,6 +60,7 @@
     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):
         """