matrix-bot/src/matrix_bot/client.py

"""
The Client class.
Connect to the matrix server and handle interactions with the server.
"""

import asyncio
import nio
from typing import (
    Optional,
    NoReturn,
    Union
)
from .async_utils import Aobject
from .utils import (
    Room,
    RoomAlias,
    RoomId
)

class Client(Aobject):
    """
        Connect to the matrix server and handle interactions with the
        server.

        allowed_rooms: dict of the rooms where the bot is allowed to connect, indexed
            by id (the name starting with '!'). If set to None, the bot connect to
            all room where it is invited.

        /!\ The client is initialized asyncronously: `client = await Client(...)`
    """

    __client: nio.AsyncClient
    __rooms_by_aliases: dict[RoomAlias, Room]
    __rooms_by_id: dict[RoomId, Room]
    allowed_rooms: Optional[dict[RoomId, Room]]

    async def __init__(
        self,
        username: str,
        homeserver: str,
        password: str,
        allowed_rooms_names: Optional[list[Union[RoomAlias, RoomId]]]=None
    ):
        """
            Initialize the Client.
            username: the username used by the bot
            homeserver: the matrix home server of the bot (expl: "https://matrix.org")
            password: the password of the user
            allowed_rooms: the list of the rooms where the bot is allowed to connect
                (given by room id (expl: '!xxx:matrix.org') of room alias (expl:
                '#xxx:matrix.org'))
        """
        self.__client = nio.AsyncClient(
            homeserver,
            username
        )
        self.__rooms_by_aliases = {}
        self.__rooms_by_id = {}
        resp = await self.__client.login(password)
        if isinstance(resp, nio.responses.LoginError):
            raise RuntimeError(f"Fail to connect: {resp.message}")

        # TODO: Where is the async map when you need it?
        self.allowed_rooms = None
        if allowed_rooms_names:
            self.allowed_rooms = {}
            for room_name in allowed_rooms_names:
                room = await self.resolve_room(room_name)
                self.allowed_rooms[room.id] = room # room uniqueness is handled by self.resolve_room


    async def resolve_room(
        self,
        room_name: Union[RoomAlias, RoomId]
    )->Room:
        """
            Lookup a room from its id or alias.
            If the name has already been resolved by this client, the
            room is return directly without querying the server.
        """
        # If the room_name is empty:
        if len(room_name) == 0:
            raise ValueError(f"Invalid room_name: {room_name}")
        # If it is a known room id:
        if room_name[0] == '!' and room_name in self.__rooms_by_id:
            return self.__rooms_by_id[room_name]
        # If it is a unknown room id:
        elif room_name[0] == '!':
            return Room(id=room_name)
        # If it is not a room id nor a room alias:
        elif room_name[0] != '#':
            raise ValueError(f"Invalid room_name: {room_name}")
        # If it is a known room alias:
        elif room_name in self.__rooms_by_aliases:
            return self.__rooms_by_aliases[room_name]
        # If it is an unknown room alias:
        else:
            resp = await self.__client.room_resolve_alias(room_name)
            if isinstance(resp, nio.responses.RoomResolveAliasError):
                raise RuntimeError(f"Error while resolving alias: {resp.message}")
            # If the room is already known:
            if resp.room_id in self.__rooms_by_id:
                room = self.__rooms_by_id[resp.room_id]
                room.aliases.add(room_name)
            # If the room is unknwon:
            else:
                room = Room(id=resp.room_id,aliases={room_name})
                self.__rooms_by_id[resp.room_id] = room
            self.__rooms_by_aliases[room_name] = room
            return room

    async def sync(
        self,
        sync_delta:int=30
    )->NoReturn:
        """
            Sync with the server every sync_delta seconds.
            sync_delta: the time in sec between each sync.
        """
        while True:
            sync_resp = await self.__client.sync(sync_delta*1000)
            if isinstance(sync_resp, nio.responses.SyncError):
                print(f"Error while syncronizing: {sync_resp.message}") # TODO: use proper logging
                continue
            print(sync_resp)

    async def run(
        self,
        sync_delta:int=30
    )->NoReturn:
        """
            Run the bot: sync with the server and execute callbacks.
        """
        await asyncio.gather(
            self.sync(sync_delta=sync_delta)
        )