Skip to content

Commit

Permalink
Migrate to asyncio
Browse files Browse the repository at this point in the history
  • Loading branch information
caibinqing committed Mar 5, 2024
1 parent fbb33ae commit aecad4e
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 20 deletions.
4 changes: 2 additions & 2 deletions custom_components/aqara_gateway/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import voluptuous as vol

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import MAJOR_VERSION, MINOR_VERSION, STATE_UNKNOWN, EVENT_HOMEASSISTANT_STOP
from homeassistant.const import MAJOR_VERSION, MINOR_VERSION, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant, Event
from homeassistant.helpers import config_validation as cv
from homeassistant.helpers.entity import Entity
Expand Down Expand Up @@ -66,7 +66,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
hass.async_create_task(hass.config_entries.async_forward_entry_setup(
entry, domain))

gateway.start()
gateway.start(hass, entry)

await hass.data[DOMAIN][entry.entry_id].async_connect()

Expand Down
47 changes: 29 additions & 18 deletions custom_components/aqara_gateway/core/gateway.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,18 @@
""" Aqara Gateway """
import asyncio
# pylint: disable=broad-except
import logging
import socket
import time
import json
import re
from threading import Thread
from typing import Optional
from random import randint
from paho.mqtt.client import Client, MQTTMessage

from homeassistant.core import Event
from homeassistant.const import CONF_NAME, CONF_PASSWORD
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import Event, HomeAssistant
from homeassistant.const import CONF_NAME, CONF_PASSWORD, MAJOR_VERSION, MINOR_VERSION
from homeassistant.components.light import ATTR_HS_COLOR, ATTR_RGB_COLOR, ATTR_BRIGHTNESS

from .shell import (
Expand All @@ -34,13 +35,14 @@
MD5_MOSQUITTO_G2HPRO_ARMV7L = '9cd591ec76f85c4d96b744eb99943eb3'
MD5_MOSQUITTO_MIPSEL = 'e0ce4757cfcccb079d89134381fd11b0'

class Gateway(Thread):
class Gateway:
# pylint: disable=too-many-instance-attributes, unused-argument
""" Aqara Gateway """

def __init__(self, hass, host: str, config: dict, **options):
main_task: asyncio.Task | None = None # for HA < 2023.3

def __init__(self, hass: HomeAssistant, host: str, config: dict, **options):
"""Initialize the Xiaomi/Aqara device."""
super().__init__(daemon=True)
self.hass = hass
self.host = host
self.options = options
Expand Down Expand Up @@ -94,9 +96,12 @@ def stop(self):
""" stop function """
self.enabled = False

async def async_connect(self) -> str:
if self.main_task: # HA < 2023.3
self.main_task.cancel()

async def async_connect(self):
"""Connect to the host. Does not process messages yet."""
result: int = None
result: int | None = None
try:
result = await self.hass.async_add_executor_job(
self._mqttc.connect,
Expand Down Expand Up @@ -125,7 +130,13 @@ def stop():

await self.hass.async_add_executor_job(stop)

def run(self):
def start(self, hass: HomeAssistant, config_entry: ConfigEntry):
if (MAJOR_VERSION, MINOR_VERSION) >= (2023, 3):
config_entry.async_create_background_task(hass, self.run(), f"{DOMAIN} gateway.run")
else:
self.main_task = hass.loop.create_task(self.run())

async def run(self):
""" Main thread loop. """
telnetshell = False
if "telnet" not in self.hass.data[DOMAIN]:
Expand All @@ -136,14 +147,14 @@ def run(self):
if not self._check_port(23):
if self.host in self.hass.data[DOMAIN]["telnet"]:
self.hass.data[DOMAIN]["telnet"].remove(self.host)
time.sleep(30)
await asyncio.sleep(30)
continue

telnetshell = True
devices = self._prepare_gateway(get_devices=True)
if isinstance(devices, list):
self._gw_topic = "gw/{}/".format(devices[0]['mac'][2:].upper())
self.setup_devices(devices)
await self.setup_devices(devices)
break

if telnetshell:
Expand All @@ -154,7 +165,7 @@ def run(self):
if not self._mqtt_connect() or not self._prepare_gateway():
if self.host in self.hass.data[DOMAIN]["mqtt"]:
self.hass.data[DOMAIN]["mqtt"].remove(self.host)
time.sleep(60)
await asyncio.sleep(60)
continue

self._mqttc.loop_start()
Expand Down Expand Up @@ -319,7 +330,7 @@ def _get_devices(self, shell):

return devices

def setup_devices(self, devices: list):
async def setup_devices(self, devices: list):
"""Add devices to hass."""
for device in devices:
timeout = 300
Expand Down Expand Up @@ -349,7 +360,7 @@ def setup_devices(self, devices: list):

# wait domain init
while domain not in self.setups and timeout > 0:
time.sleep(1)
await asyncio.sleep(1)
timeout = timeout - 1
attr = param[2]
if (attr in ('illuminance', 'light') and
Expand All @@ -360,7 +371,7 @@ def setup_devices(self, devices: list):

if self.options.get('stats'):
while 'sensor' not in self.setups:
time.sleep(1)
await asyncio.sleep(1)
self.setups['sensor'](self, device, device['type'])

def add_stats(self, ieee: str, handler):
Expand Down Expand Up @@ -467,7 +478,7 @@ def on_disconnect(self, client, userdata, ret):
self.hass.data[DOMAIN]["mqtt"].remove(self.host)
self.available = False
# self.process_gateway_stats()
self.run()
self.hass.create_task(self.run())

def on_message(self, client: Client, userdata, msg: MQTTMessage):
# pylint: disable=unused-argument
Expand Down Expand Up @@ -575,7 +586,7 @@ def _process_devices_info(self, prop, value):
'model_ver': dev['model_ver'],
'status': dev['status']
}
self.setup_devices([device])
self.hass.create_task(self.setup_devices([device]))
break

def _process_message(self, data: dict):
Expand Down Expand Up @@ -739,7 +750,7 @@ def _process_message(self, data: dict):
device['mac'] = '0x' + device['mac']
device['type'] = 'zigbee'
device['init'] = payload
self.setup_devices([device])
self.hass.create_task(self.setup_devices([device]))

async def _handle_device_remove(self, payload: dict):
"""Remove device from Hass. """
Expand Down

0 comments on commit aecad4e

Please sign in to comment.