From e6e2b32f9b956cdd386da1a5301e936059dd7743 Mon Sep 17 00:00:00 2001
From: Florian Einfalt <florian.einfalt@me.com>
Date: Tue, 28 Feb 2017 14:49:11 +0000
Subject: [PATCH] Add import_nuke function

---
 nukecontexts/__init__.py | 14 +++++++++++++-
 nukecontexts/ctx.py      |  5 +++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/nukecontexts/__init__.py b/nukecontexts/__init__.py
index 1267b9e..4df4083 100644
--- a/nukecontexts/__init__.py
+++ b/nukecontexts/__init__.py
@@ -5,6 +5,18 @@
 import logging
 import platform
 
+
+def import_nuke():
+    try:
+        import nuke
+        return nuke
+    except ImportError as e:
+        try:
+            os.environ['NON_PRODUCTION_CONTEXT']
+        except KeyError:
+            raise e
+
+
 TESTING = False
 try:
     TESTING = os.environ['NON_PRODUCTION_CONTEXT']
@@ -21,7 +33,7 @@
     match = re.search(application, sys.executable)
     if not match:
         raise RuntimeError('Import nukecontexts from within Nuke')
-    import nuke
+    nuke = import_nuke()
 
 __version__ = '0.1.4'
 __all__ = ['ctx']
diff --git a/nukecontexts/ctx.py b/nukecontexts/ctx.py
index 636208a..855b026 100644
--- a/nukecontexts/ctx.py
+++ b/nukecontexts/ctx.py
@@ -3,8 +3,9 @@
 
 from tqdm import tqdm
 
-from nukecontexts import logger
-from nukecontexts import sentry
+from nukecontexts import import_nuke, logger, sentry
+
+nuke = import_nuke()
 
 
 class NukeContextError(ValueError):