Migrate to using aioesphomeapi for the log runner to fix multiple issues (#5733)

This commit is contained in:
J. Nick Koston 2023-11-11 19:36:56 -06:00 committed by Jesse Hills
parent 91299f05f7
commit bd5905c59a
No known key found for this signature in database
GPG Key ID: BEAAE804EFD8E83A

View File

@ -1,71 +1,65 @@
from __future__ import annotations
import asyncio
import logging
from datetime import datetime
from typing import Optional
from typing import Any
from aioesphomeapi import APIClient, ReconnectLogic, APIConnectionError, LogLevel
import zeroconf
from aioesphomeapi import APIClient
from aioesphomeapi.api_pb2 import SubscribeLogsResponse
from aioesphomeapi.log_runner import async_run
from zeroconf.asyncio import AsyncZeroconf
from esphome.const import CONF_KEY, CONF_PASSWORD, CONF_PORT, __version__
from esphome.core import CORE
from esphome.const import CONF_KEY, CONF_PORT, CONF_PASSWORD, __version__
from esphome.util import safe_print
from . import CONF_ENCRYPTION
_LOGGER = logging.getLogger(__name__)
async def async_run_logs(config, address):
"""Run the logs command in the event loop."""
conf = config["api"]
port: int = int(conf[CONF_PORT])
password: str = conf[CONF_PASSWORD]
noise_psk: Optional[str] = None
noise_psk: str | None = None
if CONF_ENCRYPTION in conf:
noise_psk = conf[CONF_ENCRYPTION][CONF_KEY]
_LOGGER.info("Starting log output from %s using esphome API", address)
aiozc = AsyncZeroconf()
cli = APIClient(
address,
port,
password,
client_info=f"ESPHome Logs {__version__}",
noise_psk=noise_psk,
zeroconf_instance=aiozc.zeroconf,
)
first_connect = True
dashboard = CORE.dashboard
def on_log(msg):
time_ = datetime.now().time().strftime("[%H:%M:%S]")
text = msg.message.decode("utf8", "backslashreplace")
safe_print(time_ + text)
async def on_connect():
nonlocal first_connect
try:
await cli.subscribe_logs(
on_log,
log_level=LogLevel.LOG_LEVEL_VERY_VERBOSE,
dump_config=first_connect,
)
first_connect = False
except APIConnectionError:
cli.disconnect()
async def on_disconnect(expected_disconnect: bool) -> None:
_LOGGER.warning("Disconnected from API")
zc = zeroconf.Zeroconf()
reconnect = ReconnectLogic(
client=cli,
on_connect=on_connect,
on_disconnect=on_disconnect,
zeroconf_instance=zc,
)
await reconnect.start()
def on_log(msg: SubscribeLogsResponse) -> None:
"""Handle a new log message."""
time_ = datetime.now()
message: bytes = msg.message
text = message.decode("utf8", "backslashreplace")
if dashboard:
text = text.replace("\033", "\\033")
print(f"[{time_.hour:02}:{time_.minute:02}:{time_.second:02}]{text}")
stop = await async_run(cli, on_log, aio_zeroconf_instance=aiozc)
try:
while True:
await asyncio.sleep(60)
finally:
await aiozc.async_close()
await stop()
def run_logs(config: dict[str, Any], address: str) -> None:
"""Run the logs command."""
try:
asyncio.run(async_run_logs(config, address))
except KeyboardInterrupt:
await reconnect.stop()
zc.close()
def run_logs(config, address):
asyncio.run(async_run_logs(config, address))
pass