reserve atexit for python3 only

lbrindze · lbrindze · commit aade9746e5cc · 2016-12-02T10:59:29.000-08:00
diff --git a/analytics/client.py b/analytics/client.py
@@ -1,13 +1,13 @@
+import sys
 from datetime import datetime
 from uuid import uuid4
 import logging
 import numbers
-import atexit
 
 from dateutil.tz import tzutc
 from six import string_types
 
-from analytics.utils import guess_timezone, clean
+from analytics.utils import guess_timezone, clean, clean_exit
 from analytics.consumer import Consumer
 from analytics.version import VERSION
 
@@ -45,7 +45,7 @@ def __init__(self, write_key=None, debug=False, max_queue_size=10000,
             # destroyed before the daemon thread finishes execution. However, it
             # is *not* the same as flushing the queue! To guarantee all messages
             # have been delivered, you'll still need to call flush().
-            atexit.register(self.join)
+            clean_exit(self.join)
             self.consumer.start()
 
     def identify(self, user_id=None, traits=None, context=None, timestamp=None,
diff --git a/analytics/utils.py b/analytics/utils.py
@@ -5,6 +5,7 @@
 import numbers
 
 import six
+import sys
 
 log = logging.getLogger('segment')
 
@@ -74,3 +75,17 @@ def _coerce_unicode(cmplx):
     except:
         raise
     return item
+
+
+def clean_exit(func):
+    # ugly hack to get around
+    # atexit's excentricities
+    if (sys.version_info > (3, 0)):
+        import atexit
+        # This prevents exceptions and a messy shutdown when the interpreter is
+        # destroyed before the daemon thread finishes execution. However, it
+        # is *not* the same as flushing the queue! To guarantee all messages
+        # have been delivered, you'll still need to call flush().
+        atexit.register(func)
+    else:
+        sys.exitfunc = func
