From cd1d0543fea070425d9f088eb289268afc2221a8 Mon Sep 17 00:00:00 2001
From: InsanePrawn <insane.prawny@gmail.com>
Date: Sun, 16 Apr 2023 05:22:17 +0200
Subject: [PATCH] wrapper: move at_exit handling into wrap()

---
 wrapper/docker.py  |  1 +
 wrapper/wrapper.py | 13 +++++--------
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/wrapper/docker.py b/wrapper/docker.py
index 682b0da..52da4c1 100644
--- a/wrapper/docker.py
+++ b/wrapper/docker.py
@@ -23,6 +23,7 @@ class DockerWrapper(Wrapper):
     type: str = 'docker'
 
     def wrap(self):
+        super().wrap()
         script_path = config.runtime.script_source_dir
         assert script_path
         with open(os.path.join(script_path, 'version.txt')) as version_file:
diff --git a/wrapper/wrapper.py b/wrapper/wrapper.py
index 61bc197..35a5f2a 100644
--- a/wrapper/wrapper.py
+++ b/wrapper/wrapper.py
@@ -77,13 +77,6 @@ class Wrapper(WrapperProtocol):
     ) -> str:
         wrapped_config = f'{target_path.rstrip("/")}/{self.identifier}_wrapped.toml'
 
-        # FIXME: these at_exit hooks should go and be called from somewhere better suited
-        def at_exit():
-            self.stop()
-            os.remove(wrapped_config)
-
-        atexit.register(at_exit)
-
         dump_config_file(
             file_path=wrapped_config,
             config=(config.file | {
@@ -93,8 +86,12 @@ class Wrapper(WrapperProtocol):
         self.wrapped_config_path = wrapped_config
         return wrapped_config
 
+    def at_exit(self):
+        os.remove(self.wrapped_config_path)
+        self.stop()
+
     def wrap(self):
-        raise NotImplementedError()
+        atexit.register(self.at_exit)
 
     def stop(self):
         raise NotImplementedError()