ytdl + event pipe

disco/voice.py

@@ -1,22 +1,71 @@
+import os
+import json
+
+import gevent
+from gevent.os import make_nonblocking, nb_read
+
 from disco.gateway.packets import OPCode
 from disco.types.channel import Channel
+from disco.util.emitter import Emitter
 from telecom import TelecomConnection, AvConvPlayable
 
+try:
+    import youtube_dl
+    ytdl = youtube_dl.YoutubeDL()
+except ImportError:
+    ytdl = None
+
+
+class YoutubeDLPlayable(AvConvPlayable):
+    def __init__(self, url):
+        url = next(self.from_url(url), None)
+        if not url:
+            raise Exception('No result found for URL {}'.format(url))
+        super(YoutubeDLPlayable, self).__init__(url)
+
+    @classmethod
+    def from_url(cls, url):
+        assert ytdl is not None, 'YoutubeDL isn\'t installed'
+
+        results = ytdl.extract_info(url, download=False)
+        if 'entries' not in results:
+            results = [results]
+        else:
+            results = results['entries']
+
+        for result in results:
+            audio_formats = [fmt for fmt in result['formats'] if fmt['vcodec'] == 'none' and fmt['acodec'] == 'opus']
+            if not audio_formats:
+                raise Exception("Couldn't find valid audio format for {}".format(url))
+
+            best_audio_format = sorted(audio_formats, key=lambda i: i['abr'], reverse=True)[0]
+            yield AvConvPlayable(best_audio_format['url'])
+
 
 class VoiceConnection(object):
-    def __init__(self, client, guild_id):
+    def __init__(self, client, guild_id, enable_events=False):
         self.client = client
         self.guild_id = guild_id
         self.channel_id = None
+        self.enable_events = enable_events
         self._conn = None
         self._voice_server_update_listener = self.client.events.on(
             'VoiceServerUpdate',
             self._on_voice_server_update,
         )
+        self._event_reader_greenlet = None
+
+        self.events = None
+        if self.enable_events:
+            self.events = Emitter()
 
         self._mute = False
         self._deaf = False
 
+    def __del__(self):
+        if self._event_reader_greenlet:
+            self._event_reader_greenlet.kill()
+
     @property
     def mute(self):
         return self._mute
@@ -42,9 +91,9 @@ class VoiceConnection(object):
         self._send_voice_state_update()
 
     @classmethod
-    def from_channel(self, channel):
+    def from_channel(self, channel, **kwargs):
         assert channel.is_voice, 'Cannot connect to a non voice channel'
-        conn = VoiceConnection(channel.client, channel.guild_id)
+        conn = VoiceConnection(channel.client, channel.guild_id, **kwargs)
         conn.connect(channel.id)
         return conn
 
@@ -66,16 +115,30 @@ class VoiceConnection(object):
             self.client.gw.session_id,
         )
 
+        if self.enable_events:
+            r, w = os.pipe()
+
+            self._event_reader_greenlet = gevent.spawn(self._event_reader, r)
+            self._conn.set_event_pipe(w)
+
     def disconnect(self):
         assert self._conn is not None, 'Not connected'
 
         # Send disconnection
         self.set_channel(None)
 
+        # If we have an event reader, kill it
+        if self._event_reader_greenlet:
+            self._event_reader_greenlet.kill()
+            self._event_reader_greenlet = None
+
         # Delete our connection so it will get GC'd
         del self._conn
         self._conn = None
 
+    def play(self, playable):
+        self._conn.play(playable)
+
     def play_file(self, url):
         self._conn.play(AvConvPlayable(url))
 
@@ -93,3 +156,21 @@ class VoiceConnection(object):
             'guild_id': self.guild_id,
             'channel_id': self.channel_id,
         })
+
+    def _event_reader(self, fd):
+        if not make_nonblocking(fd):
+            raise Exception('failed to make event pipe nonblocking')
+
+        buff = ""
+        while True:
+            buff += nb_read(fd, 2048).decode('utf-8')
+
+            parts = buff.split('\n')
+            for message in parts[:-1]:
+                event = json.loads(message)
+                self.events.emit(event['e'], event['d'])
+
+            if len(parts) > 1:
+                buff = parts[-1]
+            else:
+                buff = ""

examples/music.py

@@ -1,5 +1,5 @@
 from disco.bot import Plugin
-from disco.voice import VoiceConnection
+from disco.voice import VoiceConnection, YoutubeDLPlayable
 
 
 class MusicPlugin(Plugin):
@@ -20,11 +20,13 @@ class MusicPlugin(Plugin):
                 self._connections[event.guild.id].set_channel(vs.channel)
                 return
 
-        self._connections[event.guild.id] = VoiceConnection.from_channel(vs.channel)
+        self._connections[event.guild.id] = VoiceConnection.from_channel(vs.channel, enable_events=True)
 
     @Plugin.command('play', '<song:str>')
-    def on_play(self, event, song=None):
+    def on_play(self, event, song):
         if event.guild.id not in self._connections:
             return event.msg.reply('not in voice here')
 
-        self._connections[event.guild.id].play_file(song)
+        playables = list(YoutubeDLPlayable.from_url(song))
+        for playable in playables:
+            self._connections[event.guild.id].play(playable)

setup.py

@@ -10,7 +10,7 @@ with open('README.md') as f:
     readme = f.read()
 
 extras_require = {
-    'voice': ['telecom==0.0.2'],
+    'voice': ['telecom-py==0.0.4'],
     'http': ['flask==0.12.2'],
     'yaml': ['pyyaml==3.12'],
     'music': ['youtube_dl>=2018.1.21'],