fix to use tempfile.TemporaryDirectory instead of low-level tempfile.…

…mkdtemp
HoloArchivists · Mar 31, 2022 · 7c309c0 · 7c309c0
1 parent dca4c2f
commit 7c309c0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/twspace_dl/__main__.py b/twspace_dl/__main__.py
@@ -4,7 +4,6 @@
 import json
 import logging
 import os
-import shutil
 import sys
 from typing import Iterable
 
@@ -97,8 +96,8 @@ def space(args: argparse.Namespace) -> int:
         except KeyboardInterrupt:
             logging.info("Download Interrupted")
         finally:
-            if not args.keep_files and os.path.exists(twspace_dl._tmpdir):
-                shutil.rmtree(twspace_dl._tmpdir)
+            if not args.keep_files and os.path.exists(twspace_dl.tempdir.name):
+                twspace_dl.tempdir.cleanup()
     return EXIT_CODE_SUCCESS
 
 

diff --git a/twspace_dl/twspace_dl.py b/twspace_dl/twspace_dl.py
@@ -21,7 +21,7 @@ def __init__(self, space: Twspace, format_str: str) -> None:
         self.space = space
         self.format_str = format_str or DEFAULT_FNAME_FORMAT
         self.session = requests.Session()
-        self._tmpdir: str
+        self._tempdir = tempfile.TemporaryDirectory(dir=".")
 
     @cached_property
     def filename(self) -> str:
@@ -94,13 +94,16 @@ def write_playlist(self, save_dir: str = "./") -> None:
             stream_io.write(self.playlist_text)
         logging.info(f"{path} written to disk")
 
+    @property
+    def tempdir(self):
+        return self._tempdir
+
     def download(self) -> None:
         """Download a twitter space"""
         if not shutil.which("ffmpeg"):
             raise FileNotFoundError("ffmpeg not installed")
         space = self.space
-        tempdir = self._tmpdir = tempfile.mkdtemp(dir=".")
-        self.write_playlist(save_dir=tempdir)
+        self.write_playlist(save_dir=self._tempdir.name)
         state = space["state"]
 
         cmd_base = [
@@ -121,21 +124,21 @@ def download(self) -> None:
         ]
 
         filename = os.path.basename(self.filename)
-        filename_m3u8 = os.path.join(tempdir, filename + ".m3u8")
-        filename_old = os.path.join(tempdir, filename + ".m4a")
+        filename_m3u8 = os.path.join(self._tempdir.name, filename + ".m3u8")
+        filename_old = os.path.join(self._tempdir.name, filename + ".m4a")
         cmd_old = cmd_base.copy()
         cmd_old.insert(1, "-protocol_whitelist")
         cmd_old.insert(2, "file,https,tls,tcp")
         cmd_old.insert(8, filename_m3u8)
         cmd_old.append(filename_old)
 
         if state == "Running":
-            filename_new = os.path.join(tempdir, filename + "_new.m4a")
+            filename_new = os.path.join(self._tempdir.name, filename + "_new.m4a")
             cmd_new = cmd_base.copy()
             cmd_new.insert(6, (self.dyn_url))
             cmd_new.append(filename_new)
 
-            concat_fn = os.path.join(tempdir, "list.txt")
+            concat_fn = os.path.join(self._tempdir.name, "list.txt")
             with open(concat_fn, "w", encoding="utf-8") as list_io:
                 list_io.write(
                     "file "