improve type hints

Author: theomonnom

examples/publish_hue.py

@@ -25,7 +25,7 @@ async def publish_frames(source: livekit.VideoSource):
             0, livekit.VideoRotation.VIDEO_ROTATION_0, argb_frame.to_i420())
 
         rgb = colorsys.hsv_to_rgb(hue, 1.0, 1.0)
-        rgb = [int(x * 255) for x in rgb]
+        rgb = [(x * 255) for x in rgb] # type: ignore
 
         argb_color = np.array(rgb + [255], dtype=np.uint8)
         arr.flat[::4] = argb_color[0]

examples/publish_wave.py

@@ -50,7 +50,7 @@ async def main() -> None:
         logging.info("connected to room %s", room.name)
     except livekit.ConnectError as e:
         logging.error("failed to connect to the room: %s", e)
-        return False
+        return
 
     # publish a track
     source = livekit.AudioSource()

livekit/_ffi_client.py

@@ -33,59 +33,47 @@
 
 INVALID_HANDLE = 0
 
-
-@ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t) 
-def ffi_event_callback(data_ptr: ctypes.POINTER(ctypes.c_uint8), data_len: ctypes.c_size_t) -> None: # type: ignore
+@ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t)
+def ffi_event_callback(data_ptr: ctypes.POINTER(ctypes.c_uint8), data_len: ctypes.c_size_t) -> None:  # type: ignore
     event_data = bytes(data_ptr[:int(data_len)])
     event = proto_ffi.FfiEvent()
     event.ParseFromString(event_data)
 
-    ffi_client = FfiClient()
     with ffi_client._lock:
         loop = ffi_client._event_loop
 
     loop.call_soon_threadsafe(dispatch_event, event)
 
 
 def dispatch_event(event: proto_ffi.FfiEvent) -> None:
-    ffi_client = FfiClient()
     which = str(event.WhichOneof('message'))
     ffi_client.emit(which, getattr(event, which))
 
 
-class Singleton(type):
-    _instances = {}
-
-    def __call__(cls, *args, **kwargs):
-        if cls not in cls._instances:
-            cls._instances[cls] = super(
-                Singleton, cls).__call__(*args, **kwargs)
-        return cls._instances[cls]
-
-
-class FfiClient(EventEmitter, metaclass=Singleton):
+class FfiClient(EventEmitter):
     def __init__(self) -> None:
         super().__init__()
         self._lock = threading.Lock()
 
         req = proto_ffi.FfiRequest()
-        req.initialize.event_callback_ptr = ctypes.cast(
-            ffi_event_callback, ctypes.c_void_p).value
+        cb_callback = int(ctypes.cast(
+            ffi_event_callback, ctypes.c_void_p).value) # type: ignore
+        req.initialize.event_callback_ptr = cb_callback
         self.request(req)
 
     def set_event_loop(self, loop: asyncio.AbstractEventLoop) -> None:
         with self._lock:
             self._event_loop = loop
 
     def request(self, req: proto_ffi.FfiRequest) -> proto_ffi.FfiResponse:
-        data = req.SerializeToString()
-        data_len = len(data)
-        data = (ctypes.c_ubyte * data_len)(*data)
+        proto_data = req.SerializeToString()
+        proto_len = len(proto_data)
+        data = (ctypes.c_ubyte * proto_len)(*proto_data)
 
         resp_ptr = ctypes.POINTER(ctypes.c_ubyte)()
         resp_len = ctypes.c_size_t()
         handle = ffi_lib.livekit_ffi_request(
-            data, data_len, ctypes.byref(resp_ptr), ctypes.byref(resp_len))
+            data, proto_len, ctypes.byref(resp_ptr), ctypes.byref(resp_len))
 
         resp_data = bytes(resp_ptr[:resp_len.value])
         resp = proto_ffi.FfiResponse()
@@ -94,7 +82,6 @@ def request(self, req: proto_ffi.FfiRequest) -> proto_ffi.FfiResponse:
         FfiHandle(handle)
         return resp
 
-
 class FfiHandle:
     handle = INVALID_HANDLE
 
@@ -105,3 +92,6 @@ def __del__(self):
         if self.handle != INVALID_HANDLE:
             assert ffi_lib.livekit_ffi_drop_handle(
                 ctypes.c_size_t(self.handle))
+
+
+ffi_client = FfiClient()

livekit/audio_frame.py

@@ -1,6 +1,6 @@
 import ctypes
 
-from ._ffi_client import FfiClient, FfiHandle
+from ._ffi_client import FfiHandle, ffi_client
 from ._proto import audio_frame_pb2 as proto_audio_frame
 from ._proto import ffi_pb2 as proto_ffi
 
@@ -22,7 +22,6 @@ def create(sample_rate: int, num_channels: int, samples_per_channel: int) -> 'Au
         req.alloc_audio_buffer.num_channels = num_channels
         req.alloc_audio_buffer.samples_per_channel = samples_per_channel
 
-        ffi_client = FfiClient()
         resp = ffi_client.request(req)
 
         info = resp.alloc_audio_buffer.buffer

livekit/audio_source.py

@@ -1,6 +1,6 @@
 from livekit import AudioFrame
 
-from ._ffi_client import FfiClient, FfiHandle
+from ._ffi_client import FfiHandle, ffi_client
 from ._proto import audio_frame_pb2 as proto_audio_frame
 from ._proto import ffi_pb2 as proto_ffi
 
@@ -10,7 +10,6 @@ def __init__(self) -> None:
         req = proto_ffi.FfiRequest()
         req.new_audio_source.type = proto_audio_frame.AudioSourceType.AUDIO_SOURCE_NATIVE
 
-        ffi_client = FfiClient()
         resp = ffi_client.request(req)
         self._info = resp.new_audio_source.source
         self._ffi_handle = FfiHandle(self._info.handle.id)
@@ -21,5 +20,4 @@ def capture_frame(self, frame: AudioFrame) -> None:
         req.capture_audio_frame.source_handle.id = self._ffi_handle.handle
         req.capture_audio_frame.buffer_handle.id = frame._ffi_handle.handle
 
-        ffi_client = FfiClient()
         ffi_client.request(req)

livekit/audio_stream.py

@@ -4,7 +4,7 @@
 
 from livekit import Track
 
-from ._ffi_client import FfiClient, FfiHandle
+from ._ffi_client import FfiHandle, ffi_client
 from ._proto import audio_frame_pb2 as proto_audio_frame
 from ._proto import ffi_pb2 as proto_ffi
 from .audio_frame import AudioFrame
@@ -20,18 +20,17 @@ def initalize(cls) -> None:
             return
 
         cls._initialized = True
-        ffi_client = FfiClient()
         # See VideoStream for the reason we don't use the instance method for the listener
         ffi_client.add_listener('audio_stream_event',
                                 cls._on_audio_stream_event)
 
     @classmethod
     def _on_audio_stream_event(cls, event: proto_audio_frame.AudioStreamEvent) -> None:
-        stream = cls._streams.get(event.handle.id)
-        if stream is None:
+        weak_stream = cls._streams.get(event.handle.id)
+        if weak_stream is None:
             return
 
-        stream = stream()
+        stream = weak_stream()
         if stream is None:
             return
 
@@ -51,7 +50,6 @@ def __init__(self, track: Track) -> None:
         new_audio_stream.track_handle.id = track._ffi_handle.handle
         new_audio_stream.type = proto_audio_frame.AudioStreamType.AUDIO_STREAM_NATIVE
 
-        ffi_client = FfiClient()
         resp = ffi_client.request(req)
         stream_info = resp.new_audio_stream.stream
 

livekit/participant.py

@@ -1,16 +1,21 @@
 import asyncio
 import ctypes
 import weakref
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Optional, Union
+from weakref import ref
 
 from livekit import DataPacketKind, TrackPublishOptions
 
-from ._ffi_client import FfiClient
+from ._ffi_client import ffi_client
 from ._proto import ffi_pb2 as proto_ffi
 from ._proto import participant_pb2 as proto_participant
 from ._proto import room_pb2 as proto_room
 from .track import LocalAudioTrack, LocalVideoTrack, Track
-from .track_publication import TrackPublication
+from .track_publication import (
+    LocalTrackPublication,
+    RemoteTrackPublication,
+    TrackPublication,
+)
 
 if TYPE_CHECKING:
     from livekit import Room
@@ -52,12 +57,12 @@ class LocalParticipant(Participant):
     def __init__(self, info: proto_participant.ParticipantInfo, room: weakref.ref['Room']):
         super().__init__(info)
         self._room = room
+        self.tracks: dict[str, LocalTrackPublication] = {}
 
     async def publish_data(self,
-                           # TODO(theomonnom): Allow ctypes.Array as payload?
-                           payload: bytes or str,
+                           payload: Union[bytes, str],
                            kind: DataPacketKind.ValueType = DataPacketKind.KIND_RELIABLE,
-                           destination_sids: list[str] or list['RemoteParticipant'] = []) -> None:
+                           destination_sids: Optional[Union[List[str], List['RemoteParticipant']]] = None) -> None:
 
         room = self._room()
         if room is None:
@@ -70,21 +75,24 @@ async def publish_data(self,
 
         cdata = (ctypes.c_byte * data_len)(*payload)
 
-        sids = []
-        for p in destination_sids:
-            if isinstance(p, RemoteParticipant):
-                sids.append(p.sid)
-            else:
-                sids.append(p)
+
 
         req = proto_ffi.FfiRequest()
         req.publish_data.room_handle.id = room._ffi_handle.handle
         req.publish_data.data_ptr = ctypes.addressof(cdata)
         req.publish_data.data_size = data_len
         req.publish_data.kind = kind
-        req.publish_data.destination_sids.extend(sids)
 
-        ffi_client = FfiClient()
+        if destination_sids is not None:
+            sids = []
+            for p in destination_sids:
+                if isinstance(p, RemoteParticipant):
+                    sids.append(p.sid)
+                else:
+                    sids.append(p)
+
+            req.publish_data.destination_sids.extend(sids)
+
         resp = ffi_client.request(req)
         future: asyncio.Future[proto_room.PublishDataCallback] = asyncio.Future(
         )
@@ -113,8 +121,6 @@ async def publish_track(self, track: Track, options: TrackPublishOptions) -> Tra
         req.publish_track.room_handle.id = room._ffi_handle.handle
         req.publish_track.options.CopyFrom(options)
 
-        ffi_client = FfiClient()
-
         resp = ffi_client.request(req)
 
         future: asyncio.Future[proto_room.PublishTrackCallback] = asyncio.Future(
@@ -132,7 +138,7 @@ def on_publish_callback(cb: proto_room.PublishTrackCallback):
         if cb.error:
             raise PublishTrackError(cb.error)
 
-        track_publication = TrackPublication(cb.publication)
+        track_publication = LocalTrackPublication(cb.publication, ref(self))
         track_publication.track = track
         self.tracks[track_publication.sid] = track_publication
         # TODO: Update track info
@@ -142,3 +148,4 @@ def on_publish_callback(cb: proto_room.PublishTrackCallback):
 class RemoteParticipant(Participant):
     def __init__(self, info: proto_participant.ParticipantInfo):
         super().__init__(info)
+        self.tracks: dict[str, RemoteTrackPublication] = {}

livekit/room.py

@@ -6,7 +6,7 @@
 
 from livekit import ConnectionState, TrackKind
 
-from ._ffi_client import FfiClient, FfiHandle
+from ._ffi_client import FfiHandle, ffi_client
 from ._proto import ffi_pb2 as proto_ffi
 from ._proto import participant_pb2 as proto_participant
 from ._proto import room_pb2 as proto_room
@@ -26,11 +26,9 @@ def __init__(self) -> None:
         self.participants: dict[str, RemoteParticipant] = {}
         self.connection_state = ConnectionState.CONN_DISCONNECTED
 
-        ffi_client = FfiClient()
         ffi_client.add_listener('room_event', self._on_room_event)
 
     def __del__(self):
-        ffi_client = FfiClient()
         ffi_client.remove_listener('room_event', self._on_room_event)
 
     @property
@@ -50,7 +48,6 @@ def isconnected(self) -> bool:
 
     async def connect(self, url: str, token: str) -> None:
         # TODO(theomonnom): We should be more flexible about the event loop
-        ffi_client = FfiClient()
         ffi_client.set_event_loop(asyncio.get_running_loop())
 
         req = proto_ffi.FfiRequest()
@@ -84,8 +81,6 @@ async def disconnect(self) -> None:
         if not self.isconnected():
             return
 
-        ffi_client = FfiClient()
-
         req = proto_ffi.FfiRequest()
         req.disconnect.room_handle.id = self._ffi_handle.handle
 

livekit/track.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING
 
-from ._ffi_client import FfiClient, FfiHandle
+from ._ffi_client import FfiHandle, ffi_client
 from ._proto import ffi_pb2 as proto_ffi
 from ._proto import track_pb2 as proto_track
 
@@ -47,7 +47,6 @@ def create_audio_track(name: str, source: 'AudioSource') -> 'LocalAudioTrack':
         req.create_audio_track.name = name
         req.create_audio_track.source_handle.id = source._ffi_handle.handle
 
-        ffi_client = FfiClient()
         resp = ffi_client.request(req)
         track_info = resp.create_audio_track.track
         ffi_handle = FfiHandle(track_info.handle.id)
@@ -64,7 +63,6 @@ def create_video_track(name: str, source: 'VideoSource') -> 'LocalVideoTrack':
         req.create_video_track.name = name
         req.create_video_track.source_handle.id = source._ffi_handle.handle
 
-        ffi_client = FfiClient()
         resp = ffi_client.request(req)
         track_info = resp.create_video_track.track
         ffi_handle = FfiHandle(track_info.handle.id)

livekit/track_publication.py

@@ -1,11 +1,10 @@
 import weakref
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from livekit._proto import track_pb2 as proto_track
 
-from ._ffi_client import FfiClient
+from ._ffi_client import ffi_client
 from ._proto import ffi_pb2 as proto_ffi
-from ._proto import track_pb2 as proto_track
 from .track import Track
 
 if TYPE_CHECKING:
@@ -15,7 +14,7 @@
 class TrackPublication():
     def __init__(self, info: proto_track.TrackPublicationInfo):
         self._info = info
-        self.track: Track = None
+        self.track: Optional[Track] = None
 
     @property
     def sid(self) -> str:
@@ -76,5 +75,4 @@ def set_subscribed(self, subscribed: bool):
         req.set_subscribed.participant_sid = participant.sid
         req.set_subscribed.subscribe = subscribed
 
-        ffi_client = FfiClient()
         ffi_client.request(req)