Skip to content

Commit 3e33987

Browse files
plun1331plun1331
authored
Use BytesIO for sink (#842)
1 parent 0e98bc3 commit 3e33987

File tree

10 files changed

+110
-154
lines changed

10 files changed

+110
-154
lines changed

‎discord/sinks/core.py‎

Lines changed: 9 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import sys
2828
import threading
2929
import time
30+
import io
3031
from ..types import snowflake
3132

3233
from .errors import SinkException
@@ -125,9 +126,7 @@ class AudioData:
125126
"""
126127

127128
def __init__(self, file):
128-
self.file = open(file, "ab")
129-
self.dir_path = os.path.split(file)[0]
130-
129+
self.file = file
131130
self.finished = False
132131

133132
def write(self, data):
@@ -141,16 +140,12 @@ def write(self, data):
141140
def cleanup(self):
142141
if self.finished:
143142
raise SinkException("The AudioData is already finished writing.")
144-
self.file.close()
145-
self.file = os.path.join(self.dir_path, self.file.name)
143+
self.file.seek(0)
146144
self.finished = True
147145

148146
def on_format(self, encoding):
149147
if not self.finished:
150148
raise SinkException("The AudioData is still writing.")
151-
name = os.path.split(self.file)[1]
152-
name = name.split(".")[0] + f".{encoding}"
153-
self.file = os.path.join(self.dir_path, name)
154149

155150

156151
class Sink(Filters):
@@ -169,13 +164,8 @@ class Sink(Filters):
169164
finished_callback,
170165
ctx.channel,
171166
)
172-
173-
.. versionadded:: 2.1
174167
175-
Parameters
176-
----------
177-
output_path: :class:`string`
178-
A path to where the audio files should be output.
168+
.. versionadded:: 2.1
179169
180170
Raises
181171
------
@@ -184,12 +174,11 @@ class Sink(Filters):
184174
Audio may only be formatted after recording is finished.
185175
"""
186176

187-
def __init__(self, *, output_path="", filters=None):
177+
def __init__(self, *, filters=None):
188178
if filters is None:
189179
filters = default_filters
190180
self.filters = filters
191181
Filters.__init__(self, **self.filters)
192-
self.file_path = output_path
193182
self.vc = None
194183
self.audio_data = {}
195184

@@ -200,8 +189,7 @@ def init(self, vc): # called under listen
200189
@Filters.container
201190
def write(self, data, user):
202191
if user not in self.audio_data:
203-
ssrc = self.vc.get_ssrc(user)
204-
file = os.path.join(self.file_path, f"{ssrc}.pcm")
192+
file = io.BytesIO()
205193
self.audio_data.update({user: AudioData(file)})
206194

207195
file = self.audio_data[user]
@@ -215,8 +203,8 @@ def cleanup(self):
215203

216204
def get_all_audio(self):
217205
"""Gets all audio files."""
218-
return [os.path.realpath(x.file) for x in self.audio_data.values()]
219-
206+
return [x.file for x in self.audio_data.values()]
207+
220208
def get_user_audio(self, user: snowflake.Snowflake):
221209
"""Gets the audio file(s) of one specific user."""
222-
return os.path.realpath(self.audio_data.pop(user))
210+
return os.path.realpath(self.audio_data.pop(user))

‎discord/sinks/errors.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,14 +75,16 @@ class WaveSinkError(SinkException):
7575
.. versionadded:: 2.1
7676
"""
7777

78+
7879
class M4ASinkError(SinkException):
7980
"""Exception thrown when a exception occurs with :class:`M4ASink`
8081
8182
.. versionadded:: 2.1
8283
"""
8384

85+
8486
class MKASinkError(SinkException):
8587
"""Exception thrown when a exception occurs with :class:`MKAsSink`
8688
8789
.. versionadded:: 2.1
88-
"""
90+
"""

‎discord/sinks/m4a.py‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,10 @@
2121
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
2222
DEALINGS IN THE SOFTWARE.
2323
"""
24+
import io
2425
import os
2526
import subprocess
27+
import time
2628

2729
from .core import CREATE_NO_WINDOW, Filters, Sink, default_filters
2830
from .errors import M4ASinkError
@@ -35,26 +37,20 @@ class M4ASink(Sink):
3537
3638
.. versionadded:: 2.1
3739
38-
Parameters
39-
----------
40-
output_path: :class:`string`
41-
A path to where the audio files should be output.
42-
4340
Raises
4441
------
4542
ClientException
4643
An invalid encoding type was specified.
4744
Audio may only be formatted after recording is finished.
4845
"""
4946

50-
def __init__(self, *, output_path="", filters=None):
47+
def __init__(self, *, filters=None):
5148
if filters is None:
5249
filters = default_filters
5350
self.filters = filters
5451
Filters.__init__(self, **self.filters)
5552

5653
self.encoding = "m4a"
57-
self.file_path = output_path
5854
self.vc = None
5955
self.audio_data = {}
6056

@@ -63,7 +59,7 @@ def format_audio(self, audio):
6359
raise M4ASinkError(
6460
"Audio may only be formatted after recording is finished."
6561
)
66-
m4a_file = audio.file.split(".")[0] + ".m4a"
62+
m4a_file = f"{time.time()}.tmp"
6763
args = [
6864
"ffmpeg",
6965
"-f",
@@ -73,24 +69,30 @@ def format_audio(self, audio):
7369
"-ac",
7470
"2",
7571
"-i",
76-
audio.file,
72+
"-",
73+
"-f",
74+
"ipod",
7775
m4a_file,
7876
]
79-
process = None
8077
if os.path.exists(m4a_file):
8178
os.remove(
8279
m4a_file
8380
) # process will get stuck asking whether or not to overwrite, if file already exists.
8481
try:
85-
process = subprocess.Popen(args, creationflags=CREATE_NO_WINDOW)
82+
process = subprocess.Popen(args, creationflags=CREATE_NO_WINDOW,
83+
stdin=subprocess.PIPE)
8684
except FileNotFoundError:
8785
raise M4ASinkError("ffmpeg was not found.") from None
8886
except subprocess.SubprocessError as exc:
8987
raise M4ASinkError(
9088
"Popen failed: {0.__class__.__name__}: {0}".format(exc)
9189
) from exc
9290

93-
process.wait()
91+
process.communicate(audio.file.read())
92+
93+
with open(m4a_file, "rb") as f:
94+
audio.file = io.BytesIO(f.read())
95+
audio.file.seek(0)
96+
os.remove(m4a_file)
9497

95-
os.remove(audio.file)
96-
audio.on_format(self.encoding)
98+
audio.on_format(self.encoding)

‎discord/sinks/mka.py‎

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
2222
DEALINGS IN THE SOFTWARE.
2323
"""
24+
import io
2425
import os
2526
import subprocess
2627

@@ -35,26 +36,20 @@ class MKASink(Sink):
3536
3637
.. versionadded:: 2.1
3738
38-
Parameters
39-
----------
40-
output_path: :class:`string`
41-
A path to where the audio files should be output.
42-
4339
Raises
4440
------
4541
ClientException
4642
An invalid encoding type was specified.
4743
Audio may only be formatted after recording is finished.
4844
"""
4945

50-
def __init__(self, *, output_path="", filters=None):
46+
def __init__(self, *, filters=None):
5147
if filters is None:
5248
filters = default_filters
5349
self.filters = filters
5450
Filters.__init__(self, **self.filters)
5551

5652
self.encoding = "mka"
57-
self.file_path = output_path
5853
self.vc = None
5954
self.audio_data = {}
6055

@@ -63,7 +58,6 @@ def format_audio(self, audio):
6358
raise MKASinkError(
6459
"Audio may only be formatted after recording is finished."
6560
)
66-
mka_file = audio.file.split(".")[0] + ".mka"
6761
args = [
6862
"ffmpeg",
6963
"-f",
@@ -73,24 +67,23 @@ def format_audio(self, audio):
7367
"-ac",
7468
"2",
7569
"-i",
76-
audio.file,
77-
mka_file,
70+
"-",
71+
"-f",
72+
"matroska",
73+
"pipe:1"
7874
]
79-
process = None
80-
if os.path.exists(mka_file):
81-
os.remove(
82-
mka_file
83-
) # process will get stuck asking whether or not to overwrite, if file already exists.
8475
try:
85-
process = subprocess.Popen(args, creationflags=CREATE_NO_WINDOW)
76+
process = subprocess.Popen(args, creationflags=CREATE_NO_WINDOW,
77+
stdout=subprocess.PIPE, stdin=subprocess.PIPE)
8678
except FileNotFoundError:
8779
raise MKASinkError("ffmpeg was not found.") from None
8880
except subprocess.SubprocessError as exc:
8981
raise MKASinkError(
9082
"Popen failed: {0.__class__.__name__}: {0}".format(exc)
9183
) from exc
9284

93-
process.wait()
94-
95-
os.remove(audio.file)
96-
audio.on_format(self.encoding)
85+
out = process.communicate(audio.file.read())[0]
86+
out = io.BytesIO(out)
87+
out.seek(0)
88+
audio.file = out
89+
audio.on_format(self.encoding)

‎discord/sinks/mkv.py‎

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
2222
DEALINGS IN THE SOFTWARE.
2323
"""
24+
import io
2425
import os
2526
import subprocess
2627

@@ -35,26 +36,20 @@ class MKVSink(Sink):
3536
3637
.. versionadded:: 2.1
3738
38-
Parameters
39-
----------
40-
output_path: :class:`string`
41-
A path to where the audio files should be output.
42-
4339
Raises
4440
------
4541
ClientException
4642
An invalid encoding type was specified.
4743
Audio may only be formatted after recording is finished.
4844
"""
4945

50-
def __init__(self, *, output_path="", filters=None):
46+
def __init__(self, *, filters=None):
5147
if filters is None:
5248
filters = default_filters
5349
self.filters = filters
5450
Filters.__init__(self, **self.filters)
5551

5652
self.encoding = "mkv"
57-
self.file_path = output_path
5853
self.vc = None
5954
self.audio_data = {}
6055

@@ -63,7 +58,6 @@ def format_audio(self, audio):
6358
raise MKVSinkError(
6459
"Audio may only be formatted after recording is finished."
6560
)
66-
mkv_file = audio.file.split(".")[0] + ".mkv"
6761
args = [
6862
"ffmpeg",
6963
"-f",
@@ -73,24 +67,23 @@ def format_audio(self, audio):
7367
"-ac",
7468
"2",
7569
"-i",
76-
audio.file,
77-
mkv_file,
70+
"-",
71+
"-f",
72+
"matroska",
73+
"pipe:1"
7874
]
79-
process = None
80-
if os.path.exists(mkv_file):
81-
os.remove(
82-
mkv_file
83-
) # process will get stuck asking whether or not to overwrite, if file already exists.
8475
try:
85-
process = subprocess.Popen(args, creationflags=CREATE_NO_WINDOW)
76+
process = subprocess.Popen(args, #creationflags=CREATE_NO_WINDOW,
77+
stdout=subprocess.PIPE, stdin=subprocess.PIPE)
8678
except FileNotFoundError:
8779
raise MKVSinkError("ffmpeg was not found.") from None
8880
except subprocess.SubprocessError as exc:
8981
raise MKVSinkError(
9082
"Popen failed: {0.__class__.__name__}: {0}".format(exc)
9183
) from exc
9284

93-
process.wait()
94-
95-
os.remove(audio.file)
96-
audio.on_format(self.encoding)
85+
out = process.communicate(audio.file.read())[0]
86+
out = io.BytesIO(out)
87+
out.seek(0)
88+
audio.file = out
89+
audio.on_format(self.encoding)

0 commit comments

Comments
 (0)