Skip to content

Commit bdc5f91

Browse files
authored
feat(main): separate main.py & support loguru (#14)
* chore: separate main.py * feat(log): support loguru as optional dependency * feat: correctly shutdown * fix: revert log of install_loguru
1 parent f13fe1d commit bdc5f91

File tree

20 files changed

+473
-205
lines changed

20 files changed

+473
-205
lines changed

lagrange/__init__.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
from typing import Literal, Optional
2+
import asyncio
3+
4+
from .client.client import Client as Client
5+
from .utils.log import log as log
6+
from .utils.log import install_loguru as install_loguru
7+
from .utils.sign import sign_provider
8+
from .info import InfoManager
9+
from .info.app import app_list
10+
11+
12+
class Lagrange:
13+
client: Client
14+
15+
def __init__(
16+
self,
17+
uin: int,
18+
protocol: Literal["linux", "macos", "windows"] = "linux",
19+
sign_url: Optional[str] = None,
20+
device_info_path="./device.json",
21+
signinfo_path="./sig.bin"
22+
):
23+
self.im = InfoManager(uin, device_info_path, signinfo_path)
24+
self.uin = uin
25+
self.info = app_list[protocol]
26+
self.sign = sign_provider(sign_url) if sign_url else None
27+
self.events = {}
28+
self.log = log
29+
30+
def subscribe(self, event, handler):
31+
self.events[event] = handler
32+
33+
async def login(self, client: Client):
34+
if self.im.sig_info.d2:
35+
if not await client.register():
36+
return await client.login()
37+
return True
38+
else:
39+
return await client.login()
40+
41+
async def run(self):
42+
with self.im as im:
43+
self.client = Client(
44+
self.uin,
45+
self.info,
46+
im.device,
47+
im.sig_info,
48+
self.sign,
49+
)
50+
for event, handler in self.events.items():
51+
self.client.events.subscribe(event, handler)
52+
self.client.connect()
53+
status = await self.login(self.client)
54+
if not status:
55+
log.login.error("Login failed")
56+
return
57+
await self.client.wait_closed()
58+
59+
def launch(self):
60+
try:
61+
asyncio.run(self.run())
62+
except KeyboardInterrupt:
63+
self.client._task_clear()
64+
log.root.info("Program exited by user")
65+
else:
66+
log.root.info("Program exited normally")

lagrange/client/base.py

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
from lagrange.info import AppInfo, DeviceInfo, SigInfo
99
from lagrange.utils.binary.reader import Reader
10-
from lagrange.utils.log import logger
10+
from lagrange.utils.log import log
1111
from lagrange.client.wtlogin.ntlogin import (
1212
build_ntlogin_request,
1313
parse_ntlogin_response,
@@ -58,11 +58,7 @@ def __init__(
5858
self._captcha_info = ["", "", ""] # ticket, rand_str, aid
5959

6060
self._server_push_queue: asyncio.Queue[SSOPacket] = asyncio.Queue()
61-
self._tasks: Dict[str, Optional[asyncio.Task]] = {
62-
"loop": None,
63-
"push_handle": None,
64-
"heartbeat": None,
65-
}
61+
self._tasks: Dict[str, asyncio.Task] = {}
6662
self._network = ClientNetwork(
6763
sig_info,
6864
self._server_push_queue,
@@ -98,7 +94,7 @@ def get_seq(self) -> int:
9894
self._sig.sequence += 1
9995

10096
def connect(self) -> None:
101-
if not self._tasks["loop"]:
97+
if "loop" not in self._tasks:
10298
self._tasks["loop"] = asyncio.create_task(self._network.loop())
10399
self._tasks["push_handle"] = asyncio.create_task(self._push_handle_loop())
104100
self._tasks["heartbeat"] = asyncio.create_task(self._heartbeat_task())
@@ -111,9 +107,14 @@ async def disconnect(self):
111107

112108
async def stop(self):
113109
await self.disconnect()
110+
self._task_clear()
111+
112+
def _task_clear(self):
114113
for _, task in self._tasks.items():
115-
if task:
114+
if not task.done():
116115
task.cancel()
116+
# wakeup loop if it is blocked by select() with long timeout
117+
task._loop.call_soon_threadsafe(lambda: None)
117118

118119
async def _heartbeat_task(self):
119120
err_count = 0
@@ -122,16 +123,16 @@ async def _heartbeat_task(self):
122123
if not err_count:
123124
await asyncio.sleep(self._heartbeat_interval)
124125
try:
125-
logger.network.info(
126+
log.network.info(
126127
f"{await self.sso_heartbeat(True, 5) * 1000:.2f}ms to server"
127128
)
128129
except asyncio.TimeoutError:
129130
if err_count < 3:
130-
logger.network.warning("heartbeat timeout")
131+
log.network.warning("heartbeat timeout")
131132
err_count += 1
132133
continue
133134
else:
134-
logger.network.error("too many heartbeats fail, reconnecting...")
135+
log.network.error("too many heartbeats fail, reconnecting...")
135136
self.destroy_network()
136137
err_count = 0
137138

@@ -141,10 +142,13 @@ async def _push_handle_loop(self):
141142
try:
142143
await self.push_handler(sso)
143144
except Exception as e:
144-
logger.root.error("Unhandled exception on push handler", exc_info=e)
145+
log.root.error("Unhandled exception on push handler", exc_info=e)
145146

146147
async def wait_closed(self) -> None:
147-
await self._network.wait_closed()
148+
try:
149+
await self._network.wait_closed()
150+
except asyncio.CancelledError:
151+
await self.stop()
148152

149153
@property
150154
def app_info(self) -> AppInfo:
@@ -347,7 +351,7 @@ async def qrcode_login(self, refresh_interval=5) -> bool:
347351
await asyncio.sleep(refresh_interval)
348352
ret_last = await self.get_qrcode_result()
349353
if ret_code != ret_last:
350-
logger.login.info(
354+
log.login.info(
351355
f"qrcode state changed: {ret_code.name}->{ret_last.name}"
352356
)
353357
ret_code = ret_last
@@ -395,9 +399,9 @@ async def register(self) -> bool:
395399
)
396400
if parse_register_response(response.data):
397401
self._online.set()
398-
logger.login.info("Register successful")
402+
log.login.info("Register successful")
399403
return True
400-
logger.login.error("Register failure")
404+
log.login.error("Register failure")
401405
return False
402406

403407
async def sso_heartbeat(self, calc_latency=False, timeout=10) -> float:

lagrange/client/client.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
)
3434
from lagrange.pb.service.oidb import OidbRequest, OidbResponse
3535
from lagrange.utils.binary.protobuf import proto_decode, proto_encode
36-
from lagrange.utils.log import logger
36+
from lagrange.utils.log import log
3737
from lagrange.utils.operator import timestamp
3838

3939
from qrcode.main import QRCode
@@ -65,7 +65,7 @@ def __init__(
6565
super().__init__(uin, app_info, device_info, sig_info, sign_provider, use_ipv6)
6666

6767
self._events = Events()
68-
self._highway = HighWaySession(self, logger.fork("highway"))
68+
self._highway = HighWaySession(self)
6969

7070
@property
7171
def events(self) -> Events:
@@ -100,7 +100,7 @@ async def login(self, password: str = "", qrcode_path: Optional[str] = None) ->
100100
if rsp:
101101
return True
102102
except Exception as e:
103-
logger.login.error("EasyLogin fail", exc_info=e)
103+
log.login.error("EasyLogin fail", exc_info=e)
104104

105105
if password: # TODO: PasswordLogin, WIP
106106
await self._key_exchange()
@@ -110,26 +110,26 @@ async def login(self, password: str = "", qrcode_path: Optional[str] = None) ->
110110
if rsp.successful:
111111
return await self.register()
112112
elif rsp.captcha_verify:
113-
logger.root.warning("captcha verification required")
113+
log.root.warning("captcha verification required")
114114
self.submit_login_captcha(
115115
ticket=input("ticket?->"), rand_str=input("rand_str?->")
116116
)
117117
else:
118-
logger.root.error(f"Unhandled exception raised: {rsp.name}")
118+
log.root.error(f"Unhandled exception raised: {rsp.name}")
119119
else: # QrcodeLogin
120120
ret = await self.fetch_qrcode()
121121
if isinstance(ret, int):
122-
logger.root.error(f"fetch qrcode fail: {ret}")
122+
log.root.error(f"fetch qrcode fail: {ret}")
123123
else:
124124
png, _link = ret
125125
if qrcode_path:
126-
logger.root.info(f"save qrcode to '{qrcode_path}'")
126+
log.root.info(f"save qrcode to '{qrcode_path}'")
127127
with open(qrcode_path, "wb") as f:
128128
f.write(png)
129129
else:
130130
qr = QRCode()
131131
qr.add_data(_link)
132-
logger.root.info("Please scan the qrcode below")
132+
log.root.info("Please scan the qrcode below")
133133
qr.print_ascii()
134134
if await self.qrcode_login(3):
135135
return await self.register()
@@ -149,7 +149,7 @@ async def send_oidb_svc(
149149
).data
150150
)
151151
if rsp.ret_code:
152-
logger.network.error(
152+
log.network.error(
153153
f"OidbSvc(0x{cmd:X}_{sub_cmd}) return an error: ({rsp.ret_code}){rsp.err_msg}"
154154
)
155155
return rsp

lagrange/client/event.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,22 @@
11
import asyncio
2-
from typing import TYPE_CHECKING, Callable, Coroutine, Dict, Set, Type, TypeVar
2+
from typing import TYPE_CHECKING, Any, Callable, Awaitable, Dict, Set, Type, TypeVar
33

4-
from lagrange.utils.log import logger
4+
from lagrange.utils.log import log
55

66
if TYPE_CHECKING:
77
from .events import BaseEvent
88
from .client import Client
99

1010
T = TypeVar("T", bound="BaseEvent")
11-
EVENT_HANDLER = Callable[["Client", T], Coroutine[None, None, None]]
11+
EVENT_HANDLER = Callable[["Client", T], Awaitable[Any]]
1212

1313

1414
class Events:
1515
def __init__(self):
1616
self._task_group: Set[asyncio.Task] = set()
1717
self._handle_map: Dict[Type["BaseEvent"], EVENT_HANDLER] = {}
1818

19-
def subscribe(self, event: Type["BaseEvent"], handler: EVENT_HANDLER):
19+
def subscribe(self, event: Type[T], handler: EVENT_HANDLER[T]):
2020
if event not in self._handle_map:
2121
self._handle_map[event] = handler
2222
else:
@@ -31,14 +31,14 @@ async def _task_exec(self, client: "Client", event: "BaseEvent", handler: EVENT_
3131
try:
3232
await handler(client, event)
3333
except Exception as e:
34-
logger.root.error(
34+
log.root.error(
3535
"Unhandled exception on task {}".format(event), exc_info=e
3636
)
3737

3838
def emit(self, event: "BaseEvent", client: "Client"):
3939
typ = type(event)
4040
if typ not in self._handle_map:
41-
logger.root.debug(f"Unhandled event: {event}")
41+
log.root.debug(f"Unhandled event: {event}")
4242
return
4343

4444
t = asyncio.create_task(self._task_exec(client, event, self._handle_map[typ]))

lagrange/client/highway/highway.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import asyncio
2-
import logging
32
import time
43
import uuid
54
from hashlib import md5
@@ -15,6 +14,7 @@
1514
from lagrange.utils.httpcat import HttpCat
1615
from lagrange.utils.image import decoder as decoder_img
1716
from lagrange.utils.audio import decoder as decoder_audio
17+
from lagrange.utils.log import log
1818

1919
from .encoders import (
2020
encode_audio_upload_req,
@@ -30,10 +30,8 @@
3030

3131

3232
class HighWaySession:
33-
def __init__(self, client: "Client", logger: logging.Logger):
34-
if not logger:
35-
logger = logging.getLogger(__name__)
36-
self.logger = logger
33+
def __init__(self, client: "Client"):
34+
self.logger = log.fork("highway")
3735
self._client = client
3836
self._session_sig: Optional[bytes] = None
3937
self._session_key: Optional[bytes] = None

0 commit comments

Comments
 (0)