txlyre 7 hónapja
szülő
commit
59c32b8737
1 módosított fájl, 12 hozzáadás és 5 törlés
  1. 12 5
      main.py

+ 12 - 5
main.py

@@ -17,6 +17,7 @@ try:
 except ImportError:
     logger.warning("RdRand is not available.")
 
+tmp_dir = None
 
 class TemporaryFile:
     def __init__(self, name, io, delete):
@@ -96,7 +97,7 @@ def extract_wav(path):
 def extract_video(path):
     logger.info(f"Extract video: '{path}'.")
 
-    with tempfile.TemporaryDirectory() as tmpd:
+    with tempfile.TemporaryDirectory(dir=tmp_dir) as tmpd:
         run_check(
             f"ffmpeg -hide_banner -loglevel error -y -i {path} -vf mpdecimate -r 1/1 {tmpd}/%d.bmp"
         )
@@ -137,7 +138,7 @@ def whiten(data):
 
 
 def read_video(source, duration=60):
-    tmpf = NamedTemporaryFile(suffix=".mkv", mode=None)
+    tmpf = NamedTemporaryFile(suffix=".mkv", mode=None, dir=tmp_dir)
 
     run_check(
         f"ffmpeg -hide_banner -loglevel error -y -i {source} -t {duration} -acodec copy -vcodec copy {tmpf.name}"
@@ -147,7 +148,7 @@ def read_video(source, duration=60):
 
 
 def read_audio(source, duration=60):
-    tmpf = NamedTemporaryFile(suffix=".wav", mode=None)
+    tmpf = NamedTemporaryFile(suffix=".wav", mode=None, dir=tmp_dir)
 
     run_check(
         f"ffmpeg -f alsa -i {source} -t {duration} -ar 48000 -f s16le -acodec pcm_s16le {tmpf.name}"
@@ -157,7 +158,7 @@ def read_audio(source, duration=60):
 
 
 def read_audio_video(source, duration=60):
-    tmpf = NamedTemporaryFile(suffix=".mkv", mode=None)
+    tmpf = NamedTemporaryFile(suffix=".mkv", mode=None, dir=tmp_dir)
 
     run_check(
         f"ffmpeg -hide_banner -loglevel error -y -i {source} -t {duration} -acodec copy -vcodec copy {tmpf.name}"
@@ -165,7 +166,7 @@ def read_audio_video(source, duration=60):
 
     data_a = extract_video(tmpf.name)
 
-    tmpf2 = NamedTemporaryFile(suffix=".wav", mode=None)
+    tmpf2 = NamedTemporaryFile(suffix=".wav", mode=None, dir=tmp_dir)
 
     run_check(
         f"ffmpeg -hide_banner -loglevel error -y -i {tmpf.name} -vn -ar 48000 -f s16le -acodec pcm_s16le {tmpf2.name}"
@@ -276,9 +277,15 @@ if __name__ == "__main__":
     parser.add_argument("--secret-file", type=str, default="./.secret")
     parser.add_argument("--cooldown", type=int, default=0)
     parser.add_argument("--pool-url", type=str, default="https://trng.iike.ru")
+    parser.add_argument("--tmp-dir", type=str)
 
     args = parser.parse_args()
 
+    if args.tmp_dir and os.path.isdir(args.tmp_dir):
+        tmp_dir = args.tmp_dir
+
+        logger.info(f"Changed temp-dir: '{tmp_dir}'")
+
     with open(args.secret_file, "r") as f:
         lines = f.read().strip().split("\n")
         ident = lines[0].strip()