Skip to content

Commit

Permalink
Add keyword-arguments for loading and dumping tokens to Client.
Browse files Browse the repository at this point in the history
  • Loading branch information
EvieePy committed Dec 23, 2024
1 parent 3584061 commit 0493b7f
Showing 1 changed file with 50 additions and 14 deletions.
64 changes: 50 additions & 14 deletions twitchio/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ def __init__(

self._login_called: bool = False
self._has_closed: bool = False
self._dump_tokens: bool = True

# Websockets for EventSub
self._websockets: dict[str, dict[str, Websocket]] = defaultdict(dict)
Expand Down Expand Up @@ -256,7 +257,7 @@ async def setup_hook(self) -> None:
"""
...

async def login(self, *, token: str | None = None) -> None:
async def login(self, *, token: str | None = None, load_tokens: bool = True, dump_tokens: bool = True) -> None:
"""Method to login the client and generate or store an app token.
This method is called automatically when using :meth:`~.start`.
Expand All @@ -273,11 +274,18 @@ async def login(self, *, token: str | None = None) -> None:
----------
token: str | None
An optional app token to use instead of generating one automatically.
load_tokens: bool
Optional bool which indicates whether the :class:`Client` should call :meth:`.load_tokens` during
login automatically. Defaults to ``True``.
dump_tokens: bool
Optional bool which inicates whether the :class:`Client` should call :meth:`.dump_tokens` during the
:meth:`.close` automatically. Defaults to ``True``.
"""
if self._login_called:
return

self._login_called = True
self._dump_tokens = dump_tokens

if not token:
payload: ClientCredentialsPayload = await self._http.client_credentials_token()
Expand All @@ -288,8 +296,9 @@ async def login(self, *, token: str | None = None) -> None:

self._http._app_token = token

async with self._http._token_lock:
await self.load_tokens()
if load_tokens:
async with self._http._token_lock:
await self.load_tokens()

await self.setup_hook()

Expand All @@ -299,7 +308,14 @@ async def __aenter__(self) -> Self:
async def __aexit__(self, *_: Any) -> None:
await self.close()

async def start(self, token: str | None = None, *, with_adapter: bool = True) -> None:
async def start(
self,
token: str | None = None,
*,
with_adapter: bool = True,
load_tokens: bool = True,
dump_tokens: bool = True,
) -> None:
"""|coro|
Method to login and run the `Client` asynchronously on an already running event loop.
Expand All @@ -317,7 +333,12 @@ async def start(self, token: str | None = None, *, with_adapter: bool = True) ->
An optional app token to use instead of generating one automatically.
with_adapter: bool
Whether to start and run a web adapter. Defaults to `True`. See: ... for more information.
load_tokens: bool
Optional bool which indicates whether the :class:`Client` should call :meth:`.load_tokens` during
:meth:`.login` automatically. Defaults to ``True``.
dump_tokens: bool
Optional bool which inicates whether the :class:`Client` should call :meth:`.dump_tokens` during the
:meth:`.close` automatically. Defaults to ``True``.
Examples
--------
Expand All @@ -335,7 +356,7 @@ async def main() -> None:
await client.start()
"""
self.__waiter.clear()
await self.login(token=token)
await self.login(token=token, load_tokens=load_tokens, dump_tokens=dump_tokens)

if with_adapter:
await self._adapter.run()
Expand All @@ -350,7 +371,14 @@ async def main() -> None:
self._ready_event.clear()
await self.close()

def run(self, token: str | None = None, *, with_adapter: bool = True) -> None:
def run(
self,
token: str | None = None,
*,
with_adapter: bool = True,
load_tokens: bool = True,
dump_tokens: bool = True,
) -> None:
"""Method to login the client and create a continuously running event loop.
The behaviour of this method is similar to :meth:`~.start` but instead of being used in an already running
Expand All @@ -374,7 +402,12 @@ def run(self, token: str | None = None, *, with_adapter: bool = True) -> None:
An optional app token to use instead of generating one automatically.
with_adapter: bool
Whether to start and run a web adapter. Defaults to `True`. See: ... for more information.
load_tokens: bool
Optional bool which indicates whether the :class:`Client` should call :meth:`.load_tokens` during
:meth:`.login` automatically. Defaults to ``True``.
dump_tokens: bool
Optional bool which inicates whether the :class:`Client` should call :meth:`.dump_tokens` during the
:meth:`.close` automatically. Defaults to ``True``.
Examples
--------
Expand All @@ -387,7 +420,7 @@ def run(self, token: str | None = None, *, with_adapter: bool = True) -> None:

async def run() -> None:
async with self:
await self.start(token=token, with_adapter=with_adapter)
await self.start(token=token, with_adapter=with_adapter, load_tokens=load_tokens, dump_tokens=dump_tokens)

try:
asyncio.run(run())
Expand Down Expand Up @@ -429,8 +462,9 @@ async def close(self) -> None:
for socket in sockets:
await socket.close()

async with self._http._token_lock:
await self.dump_tokens()
if self._dump_tokens:
async with self._http._token_lock:
await self.dump_tokens()

self._http.cleanup()
self.__waiter.set()
Expand Down Expand Up @@ -586,8 +620,9 @@ async def load_tokens(self, path: str | None = None, /) -> None:
.. note::
This method is always called by the client during :meth:`~.login` but **before**
:meth:`~.setup_hook`.
This method is called by the client during :meth:`~.login` but **before**
:meth:`~.setup_hook` when the ``load_tokens`` keyword-argument
is ``True`` in either, :meth:`.run`, :meth:`.start` or :meth:`.login` (Default).
You can override this method to implement your own token loading logic into the client, such as from a database.
Expand Down Expand Up @@ -627,7 +662,8 @@ async def dump_tokens(self, path: str | None = None, /) -> None:
.. note::
This method is always called by the client when it is gracefully closed.
This method is called by the client when it is gracefully closed and the ``dump_tokens`` keyword-argument
is ``True`` in either, :meth:`.run`, :meth:`.start` or :meth:`.login` (Default).
.. note::
Expand Down

0 comments on commit 0493b7f

Please sign in to comment.