Change QrLogin to reduce room for error

readthedocs/misc/v2-migration-guide.rst

@@ -984,3 +984,4 @@ sign_in no longer has phone or phone_hash (these are impl details, and now it's
 send code / sign in now only expect a single phone. resend code with new phone is send code, not resend.
 sign_up code is also now a kwarg. and no longer noop if already loggedin.
 start also mandates phone= or password= as kwarg.
+qrlogin expires has been replaced with timeout and expired for parity with tos and auth. the goal is to hide the error-prone system clock and instead use asyncio's clock. recreate was removed (just call qr_login again; parity with get_tos). class renamed to QrLogin. now must be used in a contextmgr to prevent misuse.

readthedocs/modules/custom.rst

@@ -136,7 +136,7 @@ ParticipantPermissions
     :show-inheritance:
 
 
-QRLogin
+QrLogin
 =======
 
 .. automodule:: telethon.tl.custom.qrlogin

telethon/_client/auth.py

@@ -342,10 +342,9 @@ async def send_code_request(
     return _custom.SentCode._new(result)
 
 
-async def qr_login(self: 'TelegramClient', ignored_ids: typing.List[int] = None) -> _custom.QRLogin:
+def qr_login(self: 'TelegramClient', ignored_ids: typing.List[int] = None) -> _custom.QrLogin:
-    qr_login = _custom.QRLogin(self, ignored_ids or [])
+    return _custom.QrLoginManager(self, ignored_ids)
-    await qr_login.recreate()
+
-    return qr_login
 
 async def log_out(self: 'TelegramClient') -> bool:
     try:

telethon/_client/telegramclient.py

@@ -531,7 +531,7 @@ class TelegramClient:
         """
 
     @forward_call(auth.qr_login)
-    async def qr_login(self: 'TelegramClient', ignored_ids: typing.List[int] = None) -> _custom.QRLogin:
+    def qr_login(self: 'TelegramClient', ignored_ids: typing.List[int] = None) -> _custom.QrLogin:
         """
         Initiates the QR login procedure.
 
@@ -542,15 +542,18 @@ class TelegramClient:
         whether it's the URL, using the token bytes directly, or generating
         a QR code and displaying it by other means.
 
-        See the documentation for `QRLogin` to see how to proceed after this.
+        See the documentation for `QrLogin` to see how to proceed after this.
+
+        Note that the login completes once the context manager exits,
+        not after the ``wait`` method returns.
 
         Arguments
             ignored_ids (List[`int`]):
-                List of already logged-in user IDs, to prevent logging in
+                List of already logged-in session IDs, to prevent logging in
                 twice with the same user.
 
         Returns
-            An instance of `QRLogin`.
+            An instance of `QrLogin`.
 
         Example
             .. code-block:: python
@@ -558,11 +561,16 @@ class TelegramClient:
                 def display_url_as_qr(url):
                     pass  # do whatever to show url as a qr to the user
 
-                qr_login = await client.qr_login()
+                async with client.qr_login() as qr_login:
-                display_url_as_qr(qr_login.url)
+                    display_url_as_qr(qr_login.url)
 
-                # Important! You need to wait for the login to complete!
+                    # Important! You need to wait for the login to complete!
-                await qr_login.wait()
+                    # If the context manager exits before the user logs in, the client won't be logged in.
+                    try:
+                        user = await qr_login.wait()
+                        print('Welcome,', user.first_name)
+                    except asyncio.TimeoutError:
+                        print('User did not login in time')
         """
 
     @forward_call(auth.log_out)

telethon/types/__init__.py

@@ -13,7 +13,7 @@ from ._custom import (
     InlineBuilder,
     InlineResult,
     InlineResults,
-    QRLogin,
+    QrLogin,
     ParticipantPermissions,
     Chat,
     User,

telethon/types/_custom/__init__.py

@@ -10,7 +10,7 @@ from .button import Button
 from .inlinebuilder import InlineBuilder
 from .inlineresult import InlineResult
 from .inlineresults import InlineResults
-from .qrlogin import QRLogin
+from .qrlogin import QrLoginManager, QrLogin
 from .participantpermissions import ParticipantPermissions
 from .chat import Chat
 from .user import User

telethon/types/_custom/qrlogin.py

@@ -1,30 +1,75 @@
 import asyncio
 import base64
-import datetime
+import time
+import functools
 
 from ... import _tl
 from ..._events.raw import Raw
 
 
-class QRLogin:
+class QrLoginManager:
+    def __init__(self, client, ignored_ids):
+        self._client = client
+        self._request = _tl.fn.auth.ExportLoginToken(client._api_id, client._api_hash, ignored_ids or [])
+        self._event = None
+        self._handler = None
+        self._login = None
+
+    async def __aenter__(self):
+        self._event = asyncio.Event()
+        self._handler = self._client.add_event_handler(self._callback, Raw)
+
+        try:
+            qr = await self._client(self._request)
+        except:
+            self._cleanup()
+            raise
+
+        self._login = QrLogin._new(self._client, self._request, qr, self._event)
+        return self._login
+
+    async def __aexit__(self, *args):
+        try:
+            # The logic to complete the login is in wait so the user can retrieve the logged-in user
+            await self._login.wait(timeout=0)
+            # User logged-in in time
+        except asyncio.TimeoutError:
+            pass  # User did not login in time
+        finally:
+            self._cleanup()
+
+    async def _callback(self, update):
+        if isinstance(update, _tl.UpdateLoginToken):
+            self._event.set()
+
+    def _cleanup(self):
+        # Users technically could remove all raw handlers during the procedure but it's unlikely to happen
+        self._client.remove_event_handler(self._handler)
+        self._event = None
+        self._handler = None
+        self._login = None
+
+
+class QrLogin:
     """
     QR login information.
 
     Most of the time, you will present the `url` as a QR code to the user,
     and while it's being shown, call `wait`.
     """
-    def __init__(self, client, ignored_ids):
+    def __init__(self):
-        self._client = client
+        raise TypeError('You cannot create QrLogin instances by hand!')
-        self._request = _tl.fn.auth.ExportLoginToken(
-            self._client.api_id, self._client.api_hash, ignored_ids)
-        self._resp = None
 
-    async def recreate(self):
+    @classmethod
-        """
+    def _new(cls, client, request, qr, event):
-        Generates a new token and URL for a new QR code, useful if the code
+        self = cls.__new__(cls)
-        has expired before it was imported.
+        self._client = client
-        """
+        self._request = request
-        self._resp = await self._client(self._request)
+        self._qr = qr
+        self._expiry = asyncio.get_running_loop().time() + qr.expires.timestamp() - time.time()
+        self._event = event
+        self._user = None
+        return self
 
     @property
     def token(self) -> bytes:
@@ -35,7 +80,7 @@ class QRLogin:
         :tl:`auth.importLoginToken` to log the client that originally
         requested the QR login.
         """
-        return self._resp.token
+        return self._qr.token
 
     @property
     def url(self) -> str:
@@ -54,16 +99,30 @@ class QRLogin:
 
         The URL simply consists of `token` base64-encoded.
         """
-        return 'tg://login?token={}'.format(base64.urlsafe_b64encode(self._resp.token).decode('utf-8').rstrip('='))
+        return 'tg://login?token={}'.format(base64.urlsafe_b64encode(self._qr.token).decode('utf-8').rstrip('='))
 
     @property
-    def expires(self) -> datetime.datetime:
+    def timeout(self):
         """
-        The `datetime` at which the QR code will expire.
+        How many seconds are left before `client.qr_login` should be used again.
 
-        If you want to try again, you will need to call `recreate`.
+        This value is a positive floating point number, and is monotically decreasing.
+        The value will reach zero after enough seconds have elapsed. This lets you do some work
+        and call sleep on the value and still wait just long enough.
         """
-        return self._resp.expires
+        return max(0.0, self._expiry - asyncio.get_running_loop().time())
+
+    @property
+    def expired(self):
+        """
+        Returns `True` if this instance of the QR login has expired and should be re-created.
+
+        .. code-block:: python
+
+            if qr.expired:
+                qr = await client.qr_login()
+        """
+        return asyncio.get_running_loop().time() >= self._expiry
 
     async def wait(self, timeout: float = None):
         """
@@ -71,13 +130,12 @@ class QRLogin:
         either by scanning the QR, launching the URL directly, or calling the
         import method.
 
-        This method **must** be called before the QR code is scanned, and
-        must be executing while the QR code is being scanned. Otherwise, the
-        login will not complete.
-
         Will raise `asyncio.TimeoutError` if the login doesn't complete on
         time.
 
+        Note that the login can complete even if `wait` isn't used (if the
+        context-manager is kept alive for long enough and the users logs in).
+
         Arguments
             timeout (float):
                 The timeout, in seconds, to wait before giving up. By default
@@ -85,26 +143,18 @@ class QRLogin:
                 what you want.
 
         Returns
-            On success, an instance of :tl:`User`. On failure it will raise.
+            On success, an instance of `User`. On failure it will raise.
         """
+        if self._user:
+            return self._user
+
         if timeout is None:
-            timeout = (self._resp.expires - datetime.datetime.now(tz=datetime.timezone.utc)).total_seconds()
+            timeout = self.timeout
 
-        event = asyncio.Event()
+        # Will raise timeout error if it doesn't complete quick enough,
+        # which we want to let propagate
+        await asyncio.wait_for(self._event.wait(), timeout=timeout)
 
-        async def handler(_update):
-            event.set()
-
-        self._client.add_event_handler(handler, Raw(_tl.UpdateLoginToken))
-
-        try:
-            # Will raise timeout error if it doesn't complete quick enough,
-            # which we want to let propagate
-            await asyncio.wait_for(event.wait(), timeout=timeout)
-        finally:
-            self._client.remove_event_handler(handler)
-
-        # We got here without it raising timeout error, so we can proceed
         resp = await self._client(self._request)
         if isinstance(resp, _tl.auth.LoginTokenMigrateTo):
             await self._client._switch_dc(resp.dc_id)
@@ -113,7 +163,7 @@ class QRLogin:
 
         if isinstance(resp, _tl.auth.LoginTokenSuccess):
             user = resp.authorization.user
-            self._client._on_login(user)
+            self._user = self._client._update_session_state(user)
-            return user
+            return self._user
 
-        raise TypeError('Login token response was unexpected: {}'.format(resp))
+        raise RuntimeError(f'Unexpected login token response: {resp}')