diff --git a/yaboli/__init__.py b/yaboli/__init__.py index 97b805a..d258678 100644 --- a/yaboli/__init__.py +++ b/yaboli/__init__.py @@ -14,7 +14,7 @@ from .cookiejar import * from .connection import * from .exceptions import * from .room import * -from utils import * +from .utils import * __all__ = ( connection.__all__ + diff --git a/yaboli/connection.py b/yaboli/connection.py index 3be4a0c..d56a527 100644 --- a/yaboli/connection.py +++ b/yaboli/connection.py @@ -28,7 +28,7 @@ class Connection: self._stopped = False self._pingtask = None - self._runtask = asyncio.create_task(self._run()) + self._runtask = asyncio.ensure_future(self._run()) # ... aaand the connection is started. async def send(self, ptype, data=None, await_response=True): @@ -90,8 +90,8 @@ class Connection: delay = 1 # seconds while True: try: - if self._cookiejar: - cookies = [("Cookie", cookie) for cookie in self._cookiejar.sniff()] + if self.cookiejar: + cookies = [("Cookie", cookie) for cookie in self.cookiejar.sniff()] self._ws = await websockets.connect(self.url, max_size=None, extra_headers=cookies) else: self._ws = await websockets.connect(self.url, max_size=None) @@ -106,11 +106,11 @@ class Connection: await asyncio.sleep(delay) delay *= 2 else: - if self._cookiejar: + if self.cookiejar: for set_cookie in self._ws.response_headers.get_all("Set-Cookie"): - self._cookiejar.bake(set_cookie) + self.cookiejar.bake(set_cookie) - self._pingtask = asyncio.create_task(self._ping()) + self._pingtask = asyncio.ensure_future(self._ping()) return True @@ -123,7 +123,7 @@ class Connection: - make sure the ping task has finished """ - asyncio.create_task(self.disconnect_callback()) + asyncio.ensure_future(self.disconnect_callback()) # stop ping task if self._pingtask: @@ -149,7 +149,7 @@ class Connection: """ while not self._stopped: - self._connect(self.reconnect_attempts) + await self._connect(self.reconnect_attempts) try: while True: @@ -183,13 +183,7 @@ class Connection: async def _handle_next_message(self): response = await self._ws.recv() - packet = json.loads(text) - - # Deal with pending responses - pid = packet.get("id", None) - future = self._pending_responses.pop(pid, None) - if future: - future.set_result(packet) + packet = json.loads(response) ptype = packet.get("type") data = packet.get("data", None) @@ -199,8 +193,14 @@ class Connection: else: throttled = None + # Deal with pending responses + pid = packet.get("id", None) + future = self._pending_responses.pop(pid, None) + if future: + future.set_result((ptype, data, error, throttled)) + # Pass packet onto room - asyncio.create_task(self.packet_callback(ptype, data, error, throttled)) + asyncio.ensure_future(self.packet_callback(ptype, data, error, throttled)) def _wait_for_response(self, pid): future = asyncio.Future() diff --git a/yaboli/room.py b/yaboli/room.py index 0e78eaa..efafd49 100644 --- a/yaboli/room.py +++ b/yaboli/room.py @@ -1,9 +1,14 @@ +import asyncio +import logging + +logger = logging.getLogger(__name__) + from .connection import * from .exceptions import * from .utils import * -__all__ == ["Room", "Inhabitant"] +__all__ = ["Room", "Inhabitant"] class Room: @@ -13,7 +18,7 @@ class Room: CONNECTED = 1 DISCONNECTED = 2 - EXITED = 3 + CLOSED = 3 def __init__(self, roomname, inhabitant, password=None, human=False, cookiejar=None): # TODO: Connect to room etc. @@ -38,6 +43,7 @@ class Room: self._inhabitant = inhabitant self._status = Room.DISCONNECTED + self._connected_future = asyncio.Future() # TODO: Allow for all parameters of Connection() to be specified in Room(). self._connection = Connection( @@ -48,7 +54,7 @@ class Room: ) async def exit(self): - self._status = Room.EXITED + self._status = Room.CLOSED await self._connection.stop() # ROOM COMMANDS @@ -57,16 +63,62 @@ class Room: # the command will retry once the bot has reconnected. async def get_message(self, mid): - pass + if self._status == Room.CLOSED: + raise RoomClosed() + + ptype, data, error, throttled = await self._send_while_connected( + "get-message", + id=mid + ) + + return Message.from_dict(data) async def log(self, n, before_mid=None): - pass + if self._status == Room.CLOSED: + raise RoomClosed() + + if before_mid: + ptype, data, error, throttled = await self._send_while_connected( + "log", + n=n, + before=before_mid + ) + else: + ptype, data, error, throttled = await self._send_while_connected( + "log", + n=n + ) + + return [Message.from_dict(d) for d in data.get("log")] async def nick(self, nick): - pass + if self._status == Room.CLOSED: + raise RoomClosed() + + ptype, data, error, throttled = await self._send_while_connected( + "nick", + name=nick + ) + + sid = data.get("session_id") + uid = data.get("id") + from_nick = data.get("from") + to_nick = data.get("to") + return sid, uid, from_nick, to_nick async def pm(self, uid): - pass + if self._status == Room.CLOSED: + raise RoomClosed() + + ptype, data, error, throttled = await self._send_while_connected( + "pm-initiate", + user_id=uid + ) + + # Just ignoring non-authenticated errors + pm_id = data.get("pm_id") + to_nick = data.get("to_nick") + return pm_id, to_nick async def send(self, content, parent_mid=None): """ @@ -75,13 +127,13 @@ class Room: """ if parent_mid: - data = await self._send_while_connected( + ptype, data, error, throttled = await self._send_while_connected( "send", content=content, parent=parent_mid ) else: - data = await self._send_while_connected( + ptype, data, error, throttled = await self._send_while_connected( "send", content=content ) @@ -103,6 +155,7 @@ class Room: async def _receive_packet(self, ptype, data, error, throttled): # Ignoring errors and throttling for now + logger.debug(f"Received packet of type: {ptype}") functions = { "bounce-event": self._event_bounce, #"disconnect-event": self._event_disconnect, # Not important, can ignore @@ -127,7 +180,8 @@ class Room: async def _event_bounce(self, data): if self.password is not None: try: - response = await self._connection.send("auth", type=passcode, passcode=self.password) + data = {"type": passcode, "passcode": self.password} + response = await self._connection.send("auth", data=data) rdata = response.get("data") success = rdata.get("success") if not success: @@ -140,7 +194,7 @@ class Room: async def _event_hello(self, data): self.session = Session.from_dict(data.get("session")) - self.room_is_private = = data.get("room_is_private") + self.room_is_private = data.get("room_is_private") self.version = data.get("version") self.account = data.get("account", None) self.account_has_access = data.get("account_has_access", None) @@ -178,8 +232,9 @@ class Room: async def _event_ping(self, data): try: - self._connection.send() - except exceptions.ConnectionClosed: + new_data = {"time": data.get("time")} + await self._connection.send( "ping-reply", data=new_data, await_response=False) + except ConnectionClosed: pass async def _event_pm_initiate(self, data): @@ -191,34 +246,41 @@ class Room: await self._inhabitant.pm(self, from_uid, from_nick, from_room, pm_id) async def _event_send(self, data): - pass # TODO X + message = Message.from_dict(data) + + await self._inhabitant.send(self, message) + # TODO: Figure out a way to bring fast-forwarding into this async def _event_snapshot(self, data): - # update listing + # Update listing self.listing = Listing() sessions = [Session.from_dict(d) for d in data.get("listing")] for session in sessions: self.listing.add(session) - # update (and possibly set) nick - new_nick = data.get("nick", None) - if self.session: - prev_nick = self.session.nick - if new_nick != prev_nick: - self.nick(prev_nick) - self.session.nick = new_nick - - # update more room info + # Update room info self.pm_with_nick = data.get("pm_with_nick", None), self.pm_with_user_id = data.get("pm_with_user_id", None) + # Remember old nick, because we're going to try to get it back + old_nick = self.session.nick if self.session else None + new_nick = data.get("nick", None) + self.session.nick = new_nick + + if old_nick and old_nick != new_nick: + try: + await self._connection.send("nick", data={"name": old_nick}) + except ConnectionClosed: + return # Aww, we've lost connection again + # Now, we're finally connected again! self.status = Room.CONNECTED - if not self._connected_future.done(): # should always be True, I think + if not self._connected_future.done(): # Should never be done already, I think self._connected_future.set_result(None) # Let's let the inhabitant know. + logger.debug("Letting inhabitant know") log = [Message.from_dict(m) for m in data.get("log")] await self._inhabitant.connected(self, log) @@ -244,14 +306,14 @@ class Room: # REST OF THE IMPLEMENTATION - async def _send_while_connected(*args, **kwargs): + async def _send_while_connected(self, *args, **kwargs): while True: if self._status == Room.CLOSED: raise RoomClosed() try: await self.connected() - return await self._connection.send(*args, **kwargs) + return await self._connection.send(*args, data=kwargs) except ConnectionClosed: pass # just try again @@ -263,7 +325,7 @@ class Inhabitant: # ROOM EVENTS # These functions are called by the room when something happens. -# They're launched via asyncio.create_task(), so they don't block execution of the room. +# They're launched via asyncio.ensure_future(), so they don't block execution of the room. # Just overwrite the events you need (make sure to keep the arguments the same though). async def disconnected(self, room): diff --git a/yaboli/utils.py b/yaboli/utils.py index 1fb438b..b67f48c 100644 --- a/yaboli/utils.py +++ b/yaboli/utils.py @@ -8,7 +8,7 @@ __all__ = [ "mention", "mention_reduced", "similar", "format_time", "format_time_delta", "Session", "Listing", - "Message", "Log", + "Message", ]