diff --git a/roborock/devices/device_manager.py b/roborock/devices/device_manager.py index a244cbc9..f7f7bfea 100644 --- a/roborock/devices/device_manager.py +++ b/roborock/devices/device_manager.py @@ -197,7 +197,16 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat trait = a01.create(product, channel) case DeviceVersion.B01: channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device) - trait = b01.create(channel) + model_part = product.model.split(".")[-1] + if "ss" in model_part: + raise NotImplementedError( + f"Device {device.name} has unsupported version B01_{product.model.strip('.')[-1]}" + ) + elif "sc" in model_part: + # Q7 devices start with 'sc' in their model naming. + trait = b01.q7.create(channel) + else: + raise NotImplementedError(f"Device {device.name} has unsupported B01 model: {product.model}") case _: raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}") return RoborockDevice(device, product, channel, trait) diff --git a/roborock/devices/traits/b01/__init__.py b/roborock/devices/traits/b01/__init__.py index 75c2d61f..bf6d8b23 100644 --- a/roborock/devices/traits/b01/__init__.py +++ b/roborock/devices/traits/b01/__init__.py @@ -1,30 +1,5 @@ """Traits for B01 devices.""" -from roborock import RoborockB01Methods -from roborock.devices.b01_channel import send_decoded_command -from roborock.devices.mqtt_channel import MqttChannel -from roborock.devices.traits import Trait -from roborock.roborock_message import RoborockB01Props +from .q7 import Q7PropertiesApi -__all__ = [ - "PropertiesApi", -] - - -class PropertiesApi(Trait): - """API for interacting with B01 devices.""" - - def __init__(self, channel: MqttChannel) -> None: - """Initialize the B01Props API.""" - self._channel = channel - - async def query_values(self, props: list[RoborockB01Props]) -> None: - """Query the device for the values of the given Dyad protocols.""" - await send_decoded_command( - self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params={"property": props} - ) - - -def create(channel: MqttChannel) -> PropertiesApi: - """Create traits for B01 devices.""" - return PropertiesApi(channel) +__all__ = ["Q7PropertiesApi", "q7", "q10"] diff --git a/roborock/devices/traits/b01/q10/__init__.py b/roborock/devices/traits/b01/q10/__init__.py new file mode 100644 index 00000000..b3cd30d6 --- /dev/null +++ b/roborock/devices/traits/b01/q10/__init__.py @@ -0,0 +1 @@ +"""Q10""" diff --git a/roborock/devices/traits/b01/q7/__init__.py b/roborock/devices/traits/b01/q7/__init__.py new file mode 100644 index 00000000..e63803ee --- /dev/null +++ b/roborock/devices/traits/b01/q7/__init__.py @@ -0,0 +1,31 @@ +"""Traits for Q7 B01 devices. +Potentially other devices may fall into this category in the future.""" + +from roborock.devices.b01_channel import send_decoded_command +from roborock.devices.mqtt_channel import MqttChannel +from roborock.devices.traits import Trait +from roborock.roborock_message import RoborockB01Props +from roborock.roborock_typing import RoborockB01Q7Methods + +__all__ = [ + "Q7PropertiesApi", +] + + +class Q7PropertiesApi(Trait): + """API for interacting with Q7 B01 devices.""" + + def __init__(self, channel: MqttChannel) -> None: + """Initialize the B01Props API.""" + self._channel = channel + + async def query_values(self, props: list[RoborockB01Props]) -> None: + """Query the device for the values of the given Q7 properties.""" + await send_decoded_command( + self._channel, dps=10000, command=RoborockB01Q7Methods.GET_PROP, params={"property": props} + ) + + +def create(channel: MqttChannel) -> Q7PropertiesApi: + """Create traits for B01 devices.""" + return Q7PropertiesApi(channel) diff --git a/roborock/devices/traits/traits_mixin.py b/roborock/devices/traits/traits_mixin.py index 9c1c98dd..92b9597e 100644 --- a/roborock/devices/traits/traits_mixin.py +++ b/roborock/devices/traits/traits_mixin.py @@ -31,8 +31,8 @@ class TraitsMixin: zeo: a01.ZeoApi | None = None """Zeo API, if supported.""" - b01_properties: b01.PropertiesApi | None = None - """B01 properties trait, if supported.""" + b01_q7_properties: b01.Q7PropertiesApi | None = None + """B01 Q7 properties trait, if supported.""" def __init__(self, trait: Trait) -> None: """Initialize the TraitsMixin with the given trait. diff --git a/roborock/protocols/b01_protocol.py b/roborock/protocols/b01_protocol.py index 5e60071c..27a05ecf 100644 --- a/roborock/protocols/b01_protocol.py +++ b/roborock/protocols/b01_protocol.py @@ -7,7 +7,7 @@ from Crypto.Cipher import AES from Crypto.Util.Padding import pad, unpad -from roborock import RoborockB01Methods +from roborock import RoborockB01Q7Methods from roborock.exceptions import RoborockException from roborock.roborock_message import ( RoborockMessage, @@ -18,7 +18,7 @@ _LOGGER = logging.getLogger(__name__) B01_VERSION = b"B01" -CommandType = RoborockB01Methods | str +CommandType = RoborockB01Q7Methods | str ParamsType = list | dict | int | None diff --git a/roborock/roborock_typing.py b/roborock/roborock_typing.py index 5f90dc5f..23d16cb8 100644 --- a/roborock/roborock_typing.py +++ b/roborock/roborock_typing.py @@ -271,8 +271,8 @@ class RoborockCommand(str, Enum): APP_GET_ROBOT_SETTING = "app_get_robot_setting" -class RoborockB01Methods(StrEnum): - """Methods used by the Roborock B01 model.""" +class RoborockB01Q7Methods(StrEnum): + """Methods used by the Roborock Q7 model.""" GET_PROP = "prop.get" GET_MAP_LIST = "service.get_map_list"