Skip to content

Commit

Permalink
fix: try fix script cleanup (#43)
Browse files Browse the repository at this point in the history
  • Loading branch information
qingzhuozhen committed Jan 11, 2023
1 parent df1d184 commit 8c2f0b7
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions src/amplitude/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
from amplitude.event import Revenue, BaseEvent, Identify, IdentifyEvent, GroupIdentifyEvent, EventOptions
from amplitude.plugin import AmplitudeDestinationPlugin, ContextPlugin, Plugin
from amplitude.timeline import Timeline
import atexit
import threading


class Amplitude:
Expand Down Expand Up @@ -47,6 +49,7 @@ def __init__(self, api_key: str, configuration: Config = Config()):
self.configuration.api_key = api_key
self.__timeline = Timeline()
self.__timeline.setup(self)
self._register_on_exit()
self.add(AmplitudeDestinationPlugin())
self.add(ContextPlugin())

Expand Down Expand Up @@ -167,3 +170,13 @@ def shutdown(self):
"""Shutdown the client instance, not accepting new events, flush all events in buffer"""
self.configuration.opt_out = True
self.__timeline.shutdown()

def _register_on_exit(self):
"""Internal method to clean up the client instance on main thread exit"""
if hasattr(threading, "_register_atexit"):
try:
threading._register_atexit(self.shutdown)
except Exception as e:
self.configuration.logger.warning("register for exit fail")
else:
atexit.register(self.shutdown)

0 comments on commit 8c2f0b7

Please sign in to comment.