add ffprobe check and proper environment variables

examples/readwrite.py

@@ -9,117 +9,25 @@
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        'input',
-    )
     args = parser.parse_args()
 
     # load stems
-    stems, rate = stempeg.read_stems(args.input)
-
-    # load stems,
-    # resample to 96000 Hz,
-    # use multiprocessing
-    stems, rate = stempeg.read_stems(
-        args.input,
-        sample_rate=96000,
-        multiprocess=True
-    )
+    stems, rate = stempeg.read_stems(stempeg.example_stem_path())
 
-    # --> stems now has `shape=(stem x samples x channels)``
-
-    # save stems from tensor as multi-stream mp4
-    stempeg.write_stems(
-        "test.stem.m4a",
-        stems,
-        sample_rate=96000
-    )
-
-    # save stems as dict for convenience
     stems = {
         "mix": stems[0],
         "drums": stems[1],
         "bass": stems[2],
         "other": stems[3],
         "vocals": stems[4],
     }
-    # keys will be automatically used
-
-    # from dict as files
-    stempeg.write_stems(
-        "test.stem.m4a",
-        data=stems,
-        sample_rate=96000
-    )
 
-    # `write_stems` is a preset for the following settings
-    # here the output signal is resampled to 44100 Hz and AAC codec is used
     stempeg.write_stems(
-        "test.stem.m4a",
-        stems,
-        sample_rate=96000,
-        writer=stempeg.StreamsWriter(
-            codec="aac",
-            output_sample_rate=44100,
-            bitrate="256000",
-            stem_names=['mix', 'drums', 'bass', 'other', 'vocals']
-        )
-    )
-
-    # Native Instruments compatible stems
-    stempeg.write_stems(
-        "test_traktor.stem.m4a",
-        stems,
-        sample_rate=96000,
-        writer=stempeg.NIStemsWriter(
-            stems_metadata=[
-                {"color": "#009E73", "name": "Drums"},
-                {"color": "#D55E00", "name": "Bass"},
-                {"color": "#CC79A7", "name": "Other"},
-                {"color": "#56B4E9", "name": "Vocals"}
-            ]
-        )
-    )
-
-    # lets write as multistream opus (supports only 48000 khz)
-    stempeg.write_stems(
-        "test.stem.opus",
-        stems,
-        sample_rate=96000,
-        writer=stempeg.StreamsWriter(
-            output_sample_rate=48000,
-            codec="opus"
-        )
-    )
-
-    # writing to wav requires to convert streams to multichannel
-    stempeg.write_stems(
-        "test.wav",
-        stems,
-        sample_rate=96000,
-        writer=stempeg.ChannelsWriter(
-            output_sample_rate=48000
-        )
-    )
-
-    # # stempeg also supports to load merged-multichannel streams using
-    stems, rate = stempeg.read_stems(
-        "test.wav",
-        reader=stempeg.ChannelsReader(nb_channels=2)
-    )
-
-    # mp3 does not support multiple channels,
-    # therefore we have to use `stempeg.FilesWriter`
-    # outputs are named ["output/0.mp3", "output/1.mp3"]
-    # for named files, provide a dict or use `stem_names`
-    # also apply multiprocessing
-    stempeg.write_stems(
-        ("output", ".mp3"),
+        ("output", ".flac"),
         stems,
         sample_rate=rate,
         writer=stempeg.FilesWriter(
-            multiprocess=True,
-            output_sample_rate=48000,
+            output_sample_rate=44100,
             stem_names=["mix", "drums", "bass", "other", "vocals"]
         )
     )

stempeg/__init__.py

@@ -20,7 +20,7 @@
 from .write import write_audio
 from .write import FilesWriter, StreamsWriter, ChannelsWriter, NIStemsWriter
 
-from .cmds import check_available_aac_encoders
+from .cmds import check_available_aac_encoders, ffmpeg_exists, ffprobe_exists, mp4box_exists
 
 import re
 import os

stempeg/cmds.py

@@ -1,6 +1,7 @@
 import re
 import subprocess as sp
 import logging
+import os
 
 FFMPEG_PATH = None
 FFPROBE_PATH = None
@@ -20,30 +21,53 @@ def find_cmd(cmd):
     return None
 
 
-def ffmpeg_and_ffprobe_exists():
-    global FFMPEG_PATH, FFPROBE_PATH
-    if FFMPEG_PATH is None:
+def ffmpeg_exists():
+    global FFMPEG_PATH
+    # check environment variable
+    if "FFMPEG_PATH" in os.environ:
+        env_path = os.environ["FFMPEG_PATH"]
+        FFMPEG_PATH = find_cmd(env_path)
+    else:
         FFMPEG_PATH = find_cmd("ffmpeg")
 
-    if FFPROBE_PATH is None:
+    return FFMPEG_PATH is not None
+
+
+def ffprobe_exists():
+    global FFPROBE_PATH
+    if "FFPROBE_PATH" in os.environ:
+        env_path = os.environ["FFPROBE_PATH"]
+        FFPROBE_PATH = find_cmd(env_path)
+    else:
         FFPROBE_PATH = find_cmd("ffprobe")
 
-    return FFMPEG_PATH is not None and FFPROBE_PATH is not None
+    return FFPROBE_PATH is not None
 
 
 def mp4box_exists():
     global MP4BOX_PATH
-    if MP4BOX_PATH is None:
+    if "MP4BOX_PATH" in os.environ:
+        env_path = os.environ["MP4BOX_PATH"]
+        MP4BOX_PATH = find_cmd(env_path)
+    else:
         MP4BOX_PATH = find_cmd("MP4Box")
 
     return MP4BOX_PATH is not None
 
 
-if not ffmpeg_and_ffprobe_exists():
+if not ffmpeg_exists():
+    raise RuntimeError(
+        "ffmpeg could not be found! "
+        "Please install it before using stempeg. "
+        "See: https://github.com/faroit/stempeg"
+    )
+
+
+if not ffprobe_exists():
     raise RuntimeError(
-        'ffmpeg or ffprobe could not be found! '
-        'Please install them before using stempeg. '
-        'See: https://github.com/faroit/stempeg'
+        "ffprobe could not be found! "
+        "Please install it before using stempeg. "
+        "See: https://github.com/faroit/stempeg"
     )
 
 

tests/test_read.py

@@ -1,6 +1,7 @@
 import stempeg
 import numpy as np
 import pytest
+import os
 
 
 @pytest.fixture(params=[np.float16, np.float32, np.float64])
@@ -104,3 +105,15 @@ def test_info():
     fp = stempeg.example_stem_path()
     info = stempeg.Info(fp)
     S, rate = stempeg.read_stems(fp, info=info)
+
+
+def test_cmd():
+    assert stempeg.ffmpeg_exists()
+    assert stempeg.mp4box_exists()
+    assert stempeg.ffprobe_exists()
+    os.environ["FFMPEG_PATH"] = "/path_that_does_not_exist"
+    os.environ["FFPROBE_PATH"] = "/path_that_does_not_exist"
+    os.environ["MP4BOX_PATH"] = "/path_that_does_not_exist"
+    assert not stempeg.ffprobe_exists()
+    assert not stempeg.ffmpeg_exists()
+    assert not stempeg.mp4box_exists()