diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 755f4f90..57e8d890 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,10 @@ # See https://pre-commit.com for more information # See https://pre-commit.com/hooks.html for more hooks -exclude: "CHANGELOG.md" +exclude: > + (?x)^( + CHANGELOG\.md| + roborock/map/proto/.*_pb2\.py + )$ default_stages: [ pre-commit ] repos: @@ -42,7 +46,7 @@ repos: hooks: - id: mypy exclude: cli.py - additional_dependencies: [ "types-paho-mqtt" ] + additional_dependencies: [ "types-paho-mqtt", "types-protobuf" ] - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook rev: v9.23.0 hooks: diff --git a/CHANGELOG.md b/CHANGELOG.md index 190b49ed..03aa1a75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,206 @@ +## v5.0.0 (2026-03-22) + +### Bug Fixes + +- **cli**: Make clean mode option case insensitive + ([#794](https://github.com/Python-roborock/python-roborock/pull/794), + [`7baeb65`](https://github.com/Python-roborock/python-roborock/commit/7baeb658fd80b40bfb631dc250b899e659e20210)) + +- **q10**: Restrict clean mode cli choices + ([#794](https://github.com/Python-roborock/python-roborock/pull/794), + [`7baeb65`](https://github.com/Python-roborock/python-roborock/commit/7baeb658fd80b40bfb631dc250b899e659e20210)) + +### Features + +- **api**: Again rename `YXDeviceState` enum members to have consistency with V1 state values + ([#795](https://github.com/Python-roborock/python-roborock/pull/795), + [`ad3ceea`](https://github.com/Python-roborock/python-roborock/commit/ad3ceeaea6f1b12dc51513c48476bcdf2756392b)) + +- **api**: Rename `YXWaterLevel` enum values to mirror v1 values + ([#796](https://github.com/Python-roborock/python-roborock/pull/796), + [`c630022`](https://github.com/Python-roborock/python-roborock/commit/c63002264e40beb0f6e51281afd90e10e7b731cc)) + +### Refactoring + +- **q10**: Use readable YXCleanType values with legacy aliases + ([#794](https://github.com/Python-roborock/python-roborock/pull/794), + [`7baeb65`](https://github.com/Python-roborock/python-roborock/commit/7baeb658fd80b40bfb631dc250b899e659e20210)) + +- **YXCleanType**: Move legacy values to a separate dictionary and update from_value method + ([#794](https://github.com/Python-roborock/python-roborock/pull/794), + [`7baeb65`](https://github.com/Python-roborock/python-roborock/commit/7baeb658fd80b40bfb631dc250b899e659e20210)) + +- **YXCleanType**: Remove legacy test for readable public values + ([#794](https://github.com/Python-roborock/python-roborock/pull/794), + [`7baeb65`](https://github.com/Python-roborock/python-roborock/commit/7baeb658fd80b40bfb631dc250b899e659e20210)) + +- **YXCleanType**: Simplify clean type definitions and remove legacy alias support + ([#794](https://github.com/Python-roborock/python-roborock/pull/794), + [`7baeb65`](https://github.com/Python-roborock/python-roborock/commit/7baeb658fd80b40bfb631dc250b899e659e20210)) + + +## v4.26.3 (2026-03-22) + +### Bug Fixes + +- **q10**: Add tests for Q10 status values and code mappings + ([#793](https://github.com/Python-roborock/python-roborock/pull/793), + [`dce00a2`](https://github.com/Python-roborock/python-roborock/commit/dce00a2499c3976f1cd25239bc4f81d996d51a79)) + +- **q10**: Normalize status names to canonical values + ([#793](https://github.com/Python-roborock/python-roborock/pull/793), + [`dce00a2`](https://github.com/Python-roborock/python-roborock/commit/dce00a2499c3976f1cd25239bc4f81d996d51a79)) + +- **q10**: Normalize YXDeviceState status names to canonical values + ([#793](https://github.com/Python-roborock/python-roborock/pull/793), + [`dce00a2`](https://github.com/Python-roborock/python-roborock/commit/dce00a2499c3976f1cd25239bc4f81d996d51a79)) + +- **q10**: Refactor test for canonical status names using a dictionary + ([#793](https://github.com/Python-roborock/python-roborock/pull/793), + [`dce00a2`](https://github.com/Python-roborock/python-roborock/commit/dce00a2499c3976f1cd25239bc4f81d996d51a79)) + + +## v4.26.2 (2026-03-21) + +### Bug Fixes + +- **q10**: Add missing fault field to Q10Status + ([#792](https://github.com/Python-roborock/python-roborock/pull/792), + [`66d76fc`](https://github.com/Python-roborock/python-roborock/commit/66d76fc9b3cd6d6d15f5883bfa8a22c688d9b960)) + + +## v4.26.1 (2026-03-21) + +### Bug Fixes + +- Add missing DPS fields to Q10Status and fix CLEAN_PROGRESS mapping + ([#791](https://github.com/Python-roborock/python-roborock/pull/791), + [`526da01`](https://github.com/Python-roborock/python-roborock/commit/526da01d02f6b52cab3674145273448eb602620e)) + +- Correct comment for fan level in test_status_trait_refresh + ([#791](https://github.com/Python-roborock/python-roborock/pull/791), + [`526da01`](https://github.com/Python-roborock/python-roborock/commit/526da01d02f6b52cab3674145273448eb602620e)) + +### Chores + +- Disable commitlint rules for header max length and header full stop. + ([#789](https://github.com/Python-roborock/python-roborock/pull/789), + [`711f49e`](https://github.com/Python-roborock/python-roborock/commit/711f49e9a6e4d7fc964b164c7f23265979aa166b)) + +- Increase commit header maximum length to 200 + ([#789](https://github.com/Python-roborock/python-roborock/pull/789), + [`711f49e`](https://github.com/Python-roborock/python-roborock/commit/711f49e9a6e4d7fc964b164c7f23265979aa166b)) + +- Increase commit header maximum length to 200 in commitlint configuration. + ([#789](https://github.com/Python-roborock/python-roborock/pull/789), + [`711f49e`](https://github.com/Python-roborock/python-roborock/commit/711f49e9a6e4d7fc964b164c7f23265979aa166b)) + + +## v4.26.0 (2026-03-19) + +### Features + +- Q7 Get battery level ([#790](https://github.com/Python-roborock/python-roborock/pull/790), + [`91efb3b`](https://github.com/Python-roborock/python-roborock/commit/91efb3b0be3122490c362a8e2ecc1192bb98bee6)) + + +## v4.25.0 (2026-03-16) + +### Chores + +- Apply suggestions from code review + ([#788](https://github.com/Python-roborock/python-roborock/pull/788), + [`19d7674`](https://github.com/Python-roborock/python-roborock/commit/19d7674cbf98dcf1ba591d1bf71f87b370a90a55)) + +### Features + +- Add `from_any_optional` method to `CodeMapping` for flexible enum resolution with corresponding + tests. ([#788](https://github.com/Python-roborock/python-roborock/pull/788), + [`19d7674`](https://github.com/Python-roborock/python-roborock/commit/19d7674cbf98dcf1ba591d1bf71f87b370a90a55)) + +- Add `from_any_optional` method to `RoborockModeEnum` + ([#788](https://github.com/Python-roborock/python-roborock/pull/788), + [`19d7674`](https://github.com/Python-roborock/python-roborock/commit/19d7674cbf98dcf1ba591d1bf71f87b370a90a55)) + +### Refactoring + +- Simplify B01_Q10 command parsing by removing a helper function and utilizing `from_any_optional`. + ([#788](https://github.com/Python-roborock/python-roborock/pull/788), + [`19d7674`](https://github.com/Python-roborock/python-roborock/commit/19d7674cbf98dcf1ba591d1bf71f87b370a90a55)) + + +## v4.24.0 (2026-03-16) + +### Chores + +- Fix lint. ([#787](https://github.com/Python-roborock/python-roborock/pull/787), + [`08ca9aa`](https://github.com/Python-roborock/python-roborock/commit/08ca9aa2f9dfb85f41e427d62c3ef189d3a48727)) + +- Fix MAX_PLUS enum value ([#787](https://github.com/Python-roborock/python-roborock/pull/787), + [`08ca9aa`](https://github.com/Python-roborock/python-roborock/commit/08ca9aa2f9dfb85f41e427d62c3ef189d3a48727)) + +- Rename and reorder `YXFanLevel` enum members + ([#787](https://github.com/Python-roborock/python-roborock/pull/787), + [`08ca9aa`](https://github.com/Python-roborock/python-roborock/commit/08ca9aa2f9dfb85f41e427d62c3ef189d3a48727)) + +### Documentation + +- Add docstring and alias comments to the YXFanLevel enum. + ([#787](https://github.com/Python-roborock/python-roborock/pull/787), + [`08ca9aa`](https://github.com/Python-roborock/python-roborock/commit/08ca9aa2f9dfb85f41e427d62c3ef189d3a48727)) + +### Features + +- Rename and reorder `YXFanLevel` enum members + ([#787](https://github.com/Python-roborock/python-roborock/pull/787), + [`08ca9aa`](https://github.com/Python-roborock/python-roborock/commit/08ca9aa2f9dfb85f41e427d62c3ef189d3a48727)) + + +## v4.23.0 (2026-03-16) + +### Chores + +- Remove duplicate V1TraitDataConverter + ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +- Remove unused `typing.Self` import. + ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +### Documentation + +- Clarify internal usage of V1TraitDataConverter and V1TraitMixin attributes. + ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +### Features + +- Separate trait response handling logic from refresh logic and merge + ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +- Simplify V1 trait handling ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +### Refactoring + +- Make V1TraitDataConverter an abstract base class, use a dedicated LedStatusConverter, and fix a + typo in Rooms. ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +- Remove trait update listeners and centralize data conversion into dedicated converter classes + ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + +- Standardize trait data merging to `merge_trait_values` and remove direct `_parse_response` methods + from traits. ([#783](https://github.com/Python-roborock/python-roborock/pull/783), + [`9f9c1b4`](https://github.com/Python-roborock/python-roborock/commit/9f9c1b4b9271a6a63a0dbe6afd21216b13a15648)) + + ## v4.22.0 (2026-03-14) ### Features diff --git a/commitlint.config.mjs b/commitlint.config.mjs index ad7b85d1..b5e01070 100644 --- a/commitlint.config.mjs +++ b/commitlint.config.mjs @@ -8,7 +8,11 @@ export default { // Disable the rule that enforces lowercase in subject "subject-case": [0], // 0 = disable, 1 = warn, 2 = error // Disable the rule that enforces a maximum line length in the body - "body-max-line-length": [0, "always"] + "body-max-line-length": [0, "always"], + // Disable header max length for AI-generated commits + "header-max-length": [0], + // Disable the rule that prevents periods at the end of subjects + "header-full-stop": [0] }, }; diff --git a/pyproject.toml b/pyproject.toml index c7c8d17f..863523e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "python-roborock" -version = "4.22.0" +version = "5.0.0" description = "A package to control Roborock vacuums." authors = [{ name = "humbertogontijo", email = "humbertogontijo@users.noreply.github.com" }, {name="Lash-L"}, {name="allenporter"}] requires-python = ">=3.11, <4" @@ -25,6 +25,7 @@ dependencies = [ "pycryptodomex~=3.18 ; sys_platform == 'darwin'", "paho-mqtt>=1.6.1,<3.0.0", "construct>=2.10.57,<3", + "protobuf>=5,<7", "vacuum-map-parser-roborock", "pyrate-limiter>=4.0.0,<5", "aiomqtt>=2.5.0,<3", @@ -97,9 +98,15 @@ major_tags= ["refactor"] lint.ignore = ["F403", "E741"] lint.select=["E", "F", "UP", "I"] line-length = 120 +extend-exclude = ["roborock/map/proto/*_pb2.py"] [tool.ruff.lint.per-file-ignores] "*/__init__.py" = ["F401"] +"roborock/map/proto/*_pb2.py" = ["E501", "I001", "UP009"] + +[[tool.mypy.overrides]] +module = ["roborock.map.proto.*"] +ignore_errors = true [tool.pytest.ini_options] asyncio_mode = "auto" diff --git a/roborock/broadcast_protocol.py b/roborock/broadcast_protocol.py index a08bca94..e226a9cf 100644 --- a/roborock/broadcast_protocol.py +++ b/roborock/broadcast_protocol.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import asyncio import hashlib import json diff --git a/roborock/cli.py b/roborock/cli.py index 8e53f5d1..b2b208ff 100644 --- a/roborock/cli.py +++ b/roborock/cli.py @@ -764,21 +764,6 @@ async def network_info(ctx, device_id: str): await _display_v1_trait(context, device_id, lambda v1: v1.network_info) -def _parse_b01_q10_command(cmd: str) -> B01_Q10_DP: - """Parse B01_Q10 command from either enum name or value.""" - try: - return B01_Q10_DP(int(cmd)) - except ValueError: - try: - return B01_Q10_DP.from_name(cmd) - except ValueError: - try: - return B01_Q10_DP.from_value(cmd) - except ValueError: - pass - raise RoborockException(f"Invalid command {cmd} for B01_Q10 device") - - @click.command() @click.option("--device_id", required=True) @click.option("--cmd", required=True) @@ -795,7 +780,8 @@ async def command(ctx, cmd, device_id, params): if result: click.echo(dump_json(result)) elif device.b01_q10_properties is not None: - cmd_value = _parse_b01_q10_command(cmd) + if cmd_value := B01_Q10_DP.from_any_optional(cmd) is None: + raise RoborockException(f"Invalid command {cmd} for B01_Q10 device") command_trait: Trait = device.b01_q10_properties.command await command_trait.send(cmd_value, json.loads(params) if params is not None else None) click.echo("Command sent successfully; Enable debug logging (-d) to see responses.") @@ -1289,7 +1275,12 @@ async def q10_empty_dustbin(ctx: click.Context, device_id: str) -> None: @session.command() @click.option("--device_id", required=True, help="Device ID") -@click.option("--mode", required=True, type=click.Choice(["bothwork", "onlysweep", "onlymop"]), help="Clean mode") +@click.option( + "--mode", + required=True, + type=click.Choice(["vac_and_mop", "vacuum", "mop"], case_sensitive=False), + help='Clean mode (preferred: "vac_and_mop", "vacuum", "mop")', +) @click.pass_context @async_command async def q10_set_clean_mode(ctx: click.Context, device_id: str, mode: str) -> None: diff --git a/roborock/data/b01_q10/b01_q10_code_mappings.py b/roborock/data/b01_q10/b01_q10_code_mappings.py index abadb9de..8020894d 100644 --- a/roborock/data/b01_q10/b01_q10_code_mappings.py +++ b/roborock/data/b01_q10/b01_q10_code_mappings.py @@ -119,20 +119,26 @@ class B01_Q10_DP(RoborockModeEnum): class YXFanLevel(RoborockModeEnum): + """The fan or vacuum level of the robot. + + Note: The names used here are the v1 names, though the values + have different aliases in the app bundles. + """ + UNKNOWN = "unknown", -1 - CLOSE = "close", 0 + OFF = "off", 0 # close QUIET = "quiet", 1 - NORMAL = "normal", 2 - STRONG = "strong", 3 + BALANCED = "balanced", 2 # normal + TURBO = "turbo", 3 # strong MAX = "max", 4 - SUPER = "super", 8 + MAX_PLUS = "max_plus", 8 # super class YXWaterLevel(RoborockModeEnum): UNKNOWN = "unknown", -1 - CLOSE = "close", 0 + OFF = "off", 0 # close LOW = "low", 1 - MIDDLE = "middle", 2 + MEDIUM = "medium", 2 # middle HIGH = "high", 3 @@ -151,31 +157,31 @@ class YXRoomMaterial(RoborockModeEnum): class YXCleanType(RoborockModeEnum): UNKNOWN = "unknown", -1 - BOTH_WORK = "bothwork", 1 - ONLY_SWEEP = "onlysweep", 2 - ONLY_MOP = "onlymop", 3 + VAC_AND_MOP = "vac_and_mop", 1 # bothwork + VACUUM = "vacuum", 2 # onlysweep + MOP = "mop", 3 # onlymop class YXDeviceState(RoborockModeEnum): UNKNOWN = "unknown", -1 - SLEEP_STATE = "sleepstate", 2 - STANDBY_STATE = "standbystate", 3 - CLEANING_STATE = "cleaningstate", 5 - TO_CHARGE_STATE = "tochargestate", 6 - REMOTEING_STATE = "remoteingstate", 7 - CHARGING_STATE = "chargingstate", 8 - PAUSE_STATE = "pausestate", 10 - FAULT_STATE = "faultstate", 12 - UPGRADE_STATE = "upgradestate", 14 - DUSTING = "dusting", 22 - CREATING_MAP_STATE = "creatingmapstate", 29 - MAP_SAVE_STATE = "mapsavestate", 99 - RE_LOCATION_STATE = "relocationstate", 101 - ROBOT_SWEEPING = "robotsweeping", 102 - ROBOT_MOPING = "robotmoping", 103 - ROBOT_SWEEP_AND_MOPING = "robotsweepandmoping", 104 - ROBOT_TRANSITIONING = "robottransitioning", 105 - ROBOT_WAIT_CHARGE = "robotwaitcharge", 108 + SLEEPING = "sleeping", 2 # sleepstate + IDLE = "idle", 3 # standbystate + CLEANING = "cleaning", 5 # cleaningstate + RETURNING_HOME = "returning_home", 6 # tochargestate + REMOTE_CONTROL_ACTIVE = "remote_control_active", 7 # remoteingstate + CHARGING = "charging", 8 # chargingstate + PAUSED = "paused", 10 # pausestate + ERROR = "error", 12 # faultstate + UPDATING = "updating", 14 # upgradestate + EMPTYING_THE_BIN = "emptying_the_bin", 22 # dusting + MAPPING = "mapping", 29 # creatingmapstate + SAVING_MAP = "saving_map", 99 # mapsavestate + RELOCATING = "relocating", 101 # relocationstate + SWEEPING = "sweeping", 102 # robotsweeping + MOPPING = "mopping", 103 # robotmoping + SWEEP_AND_MOP = "sweep_and_mop", 104 # robotsweepandmoping + TRANSITIONING = "transitioning", 105 # robottransitioning + WAITING_TO_CHARGE = "waiting_to_charge", 108 # robotwaitcharge class YXBackType(RoborockModeEnum): diff --git a/roborock/data/b01_q10/b01_q10_containers.py b/roborock/data/b01_q10/b01_q10_containers.py index 562c7062..4b3fb3ce 100644 --- a/roborock/data/b01_q10/b01_q10_containers.py +++ b/roborock/data/b01_q10/b01_q10_containers.py @@ -96,7 +96,15 @@ class Q10Status(RoborockBase): fan_level: YXFanLevel | None = field(default=None, metadata={"dps": B01_Q10_DP.FAN_LEVEL}) water_level: YXWaterLevel | None = field(default=None, metadata={"dps": B01_Q10_DP.WATER_LEVEL}) clean_count: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_COUNT}) + total_clean_area: int | None = field(default=None, metadata={"dps": B01_Q10_DP.TOTAL_CLEAN_AREA}) + total_clean_count: int | None = field(default=None, metadata={"dps": B01_Q10_DP.TOTAL_CLEAN_COUNT}) + total_clean_time: int | None = field(default=None, metadata={"dps": B01_Q10_DP.TOTAL_CLEAN_TIME}) + main_brush_life: int | None = field(default=None, metadata={"dps": B01_Q10_DP.MAIN_BRUSH_LIFE}) + side_brush_life: int | None = field(default=None, metadata={"dps": B01_Q10_DP.SIDE_BRUSH_LIFE}) + filter_life: int | None = field(default=None, metadata={"dps": B01_Q10_DP.FILTER_LIFE}) + sensor_life: int | None = field(default=None, metadata={"dps": B01_Q10_DP.SENSOR_LIFE}) clean_mode: YXDeviceWorkMode | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_MODE}) clean_task_type: YXDeviceCleanTask | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_TASK_TYPE}) back_type: YXBackType | None = field(default=None, metadata={"dps": B01_Q10_DP.BACK_TYPE}) - cleaning_progress: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEANING_PROGRESS}) + cleaning_progress: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_PROGRESS}) + fault: int | None = field(default=None, metadata={"dps": B01_Q10_DP.FAULT}) diff --git a/roborock/data/b01_q7/b01_q7_containers.py b/roborock/data/b01_q7/b01_q7_containers.py index 6ae43099..a75b7185 100644 --- a/roborock/data/b01_q7/b01_q7_containers.py +++ b/roborock/data/b01_q7/b01_q7_containers.py @@ -114,7 +114,7 @@ class B01Props(RoborockBase): wind: SCWindMapping | None = None water: WaterLevelMapping | None = None mode: CleanTypeMapping | None = None - quantity: int | None = None + quantity: int | None = None # The Q7 L5 reports its battery level as 'quantity' alarm: int | None = None volume: int | None = None hypa: int | None = None @@ -169,6 +169,13 @@ class B01Props(RoborockBase): recommend: Recommend | None = None add_sweep_status: int | None = None + @property + def battery(self) -> int | None: + """ + Returns device battery level as a percentage. + """ + return self.quantity + @property def main_brush_time_left(self) -> int | None: """ diff --git a/roborock/data/code_mappings.py b/roborock/data/code_mappings.py index fc34b387..d1f36b83 100644 --- a/roborock/data/code_mappings.py +++ b/roborock/data/code_mappings.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import logging from collections import namedtuple from enum import Enum, IntEnum, StrEnum @@ -17,7 +15,7 @@ def name(self) -> str: return super().name.lower() @classmethod - def _missing_(cls: type[RoborockEnum], key) -> RoborockEnum: + def _missing_(cls: type[Self], key) -> Self: if hasattr(cls, "unknown"): warning = f"Missing {cls.__name__} code: {key} - defaulting to 'unknown'" if warning not in completed_warnings: @@ -32,23 +30,23 @@ def _missing_(cls: type[RoborockEnum], key) -> RoborockEnum: return default_value @classmethod - def as_dict(cls: type[RoborockEnum]): + def as_dict(cls: type[Self]): return {i.name: i.value for i in cls if i.name != "missing"} @classmethod - def as_enum_dict(cls: type[RoborockEnum]): + def as_enum_dict(cls: type[Self]): return {i.value: i for i in cls if i.name != "missing"} @classmethod - def values(cls: type[RoborockEnum]) -> list[int]: + def values(cls: type[Self]) -> list[int]: return list(cls.as_dict().values()) @classmethod - def keys(cls: type[RoborockEnum]) -> list[str]: + def keys(cls: type[Self]) -> list[str]: return list(cls.as_dict().keys()) @classmethod - def items(cls: type[RoborockEnum]): + def items(cls: type[Self]): return cls.as_dict().items() @@ -100,6 +98,30 @@ def from_name(cls, name: str) -> Self: return member raise ValueError(f"{name} is not a valid name for {cls.__name__}") + @classmethod + def from_any_optional(cls, value: str | int) -> Self | None: + """Resolve a string or int to an enum member. + + Tries to look up by enum name, string value, or integer code + and returns None if no match is found. + """ + # Try enum name lookup (e.g. "SEEK") + try: + return cls.from_name(str(value)) + except ValueError: + pass + # Try DP string value lookup (e.g. "dpSeek") + try: + return cls.from_value(str(value)) + except ValueError: + pass + # Try integer code lookup (e.g. "11") + try: + return cls.from_code(int(value)) + except (ValueError, TypeError): + pass + return None + @classmethod def keys(cls) -> list[str]: """Returns a list of all member values.""" diff --git a/roborock/data/v1/v1_code_mappings.py b/roborock/data/v1/v1_code_mappings.py index c0bab14b..2ad3a0b3 100644 --- a/roborock/data/v1/v1_code_mappings.py +++ b/roborock/data/v1/v1_code_mappings.py @@ -1,3 +1,5 @@ +from typing import Self + from ..code_mappings import RoborockEnum @@ -91,7 +93,7 @@ class RoborockStartType(RoborockEnum): class RoborockDssCodes(RoborockEnum): @classmethod - def _missing_(cls: type[RoborockEnum], key) -> RoborockEnum: + def _missing_(cls: type[Self], key) -> Self: # If the calculated value is not provided, then it should be viewed as okay. # As the math will sometimes result in you getting numbers that don't matter. return cls.okay # type: ignore diff --git a/roborock/device_features.py b/roborock/device_features.py index b8ecf483..ef3ea19f 100644 --- a/roborock/device_features.py +++ b/roborock/device_features.py @@ -1,8 +1,6 @@ -from __future__ import annotations - from dataclasses import dataclass, field, fields from enum import IntEnum, StrEnum -from typing import Any +from typing import Any, Self from roborock.data.code_mappings import RoborockProductNickname from roborock.data.containers import RoborockBase @@ -566,7 +564,7 @@ def from_feature_flags( new_feature_info_str: str, feature_info: list[int], product_nickname: RoborockProductNickname | None, - ) -> DeviceFeatures: + ) -> Self: """Creates a DeviceFeatures instance from raw feature flags. :param new_feature_info: A int from get_init_status (sometimes can be found in homedata, but it is not always) :param new_feature_info_str: A hex string from get_init_status or home_data. diff --git a/roborock/devices/device_manager.py b/roborock/devices/device_manager.py index b1ef6626..0be98ea1 100644 --- a/roborock/devices/device_manager.py +++ b/roborock/devices/device_manager.py @@ -251,7 +251,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat trait = b01.q10.create(channel) elif "sc" in model_part: # Q7 devices start with 'sc' in their model naming. - trait = b01.q7.create(channel) + trait = b01.q7.create(product, device, channel) else: raise UnsupportedDeviceError(f"Device {device.name} has unsupported B01 model: {product.model}") case _: diff --git a/roborock/devices/rpc/b01_q10_channel.py b/roborock/devices/rpc/b01_q10_channel.py index d27b148b..1e0510ba 100644 --- a/roborock/devices/rpc/b01_q10_channel.py +++ b/roborock/devices/rpc/b01_q10_channel.py @@ -1,7 +1,5 @@ """Thin wrapper around the MQTT channel for Roborock B01 Q10 devices.""" -from __future__ import annotations - import logging from collections.abc import AsyncGenerator from typing import Any diff --git a/roborock/devices/traits/b01/q7/__init__.py b/roborock/devices/traits/b01/q7/__init__.py index 9bbe6818..f29f287b 100644 --- a/roborock/devices/traits/b01/q7/__init__.py +++ b/roborock/devices/traits/b01/q7/__init__.py @@ -1,10 +1,14 @@ """Traits for Q7 B01 devices. -Potentially other devices may fall into this category in the future.""" + +Potentially other devices may fall into this category in the future. +""" + +from __future__ import annotations from typing import Any from roborock import B01Props -from roborock.data import Q7MapList, Q7MapListEntry +from roborock.data import HomeDataDevice, HomeDataProduct, Q7MapList, Q7MapListEntry from roborock.data.b01_q7.b01_q7_code_mappings import ( CleanPathPreferenceMapping, CleanRepeatMapping, @@ -23,18 +27,20 @@ from .clean_summary import CleanSummaryTrait from .map import MapTrait +from .map_content import MapContentTrait __all__ = [ "Q7PropertiesApi", "CleanSummaryTrait", "MapTrait", + "MapContentTrait", "Q7MapList", "Q7MapListEntry", ] class Q7PropertiesApi(Trait): - """API for interacting with B01 devices.""" + """API for interacting with B01 Q7 devices.""" clean_summary: CleanSummaryTrait """Trait for clean records / clean summary (Q7 `service.get_record_list`).""" @@ -42,11 +48,25 @@ class Q7PropertiesApi(Trait): map: MapTrait """Trait for map list metadata + raw map payload retrieval.""" - def __init__(self, channel: MqttChannel) -> None: - """Initialize the B01Props API.""" + map_content: MapContentTrait + """Trait for fetching parsed current map content.""" + + def __init__(self, channel: MqttChannel, *, device: HomeDataDevice, product: HomeDataProduct) -> None: + """Initialize the Q7 API.""" self._channel = channel + self._device = device + self._product = product + + if not device.sn or not product.model: + raise ValueError("B01 Q7 map content requires device serial number and product model metadata") + self.clean_summary = CleanSummaryTrait(channel) self.map = MapTrait(channel) + self.map_content = MapContentTrait( + self.map, + serial=device.sn, + model=product.model, + ) async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None: """Query the device for the values of the given Q7 properties.""" @@ -151,6 +171,6 @@ async def send(self, command: CommandType, params: ParamsType) -> Any: ) -def create(channel: MqttChannel) -> Q7PropertiesApi: - """Create traits for B01 devices.""" - return Q7PropertiesApi(channel) +def create(product: HomeDataProduct, device: HomeDataDevice, channel: MqttChannel) -> Q7PropertiesApi: + """Create traits for B01 Q7 devices.""" + return Q7PropertiesApi(channel, device=device, product=product) diff --git a/roborock/devices/traits/b01/q7/clean_summary.py b/roborock/devices/traits/b01/q7/clean_summary.py index b49040f3..65fea0e8 100644 --- a/roborock/devices/traits/b01/q7/clean_summary.py +++ b/roborock/devices/traits/b01/q7/clean_summary.py @@ -4,8 +4,6 @@ and a `record_list` whose items contain a JSON string in `detail`. """ -from __future__ import annotations - import logging from roborock import CleanRecordDetail, CleanRecordList, CleanRecordSummary diff --git a/roborock/devices/traits/b01/q7/map_content.py b/roborock/devices/traits/b01/q7/map_content.py new file mode 100644 index 00000000..db00119b --- /dev/null +++ b/roborock/devices/traits/b01/q7/map_content.py @@ -0,0 +1,98 @@ +"""Trait for fetching parsed map content from B01/Q7 devices. + +This intentionally mirrors the v1 `MapContentTrait` contract: +- `refresh()` performs I/O and populates cached fields +- `parse_map_content()` reparses cached raw bytes without I/O +- fields `image_content`, `map_data`, and `raw_api_response` are then readable + +For B01/Q7 devices, the underlying raw map payload is retrieved via `MapTrait`. +""" + +from dataclasses import dataclass + +from vacuum_map_parser_base.map_data import MapData + +from roborock.data import RoborockBase +from roborock.devices.traits import Trait +from roborock.exceptions import RoborockException +from roborock.map.b01_map_parser import B01MapParser, B01MapParserConfig + +from .map import MapTrait + +_TRUNCATE_LENGTH = 20 + + +@dataclass +class MapContent(RoborockBase): + """Dataclass representing map content.""" + + image_content: bytes | None = None + """The rendered image of the map in PNG format.""" + + map_data: MapData | None = None + """Parsed map data (metadata for points on the map).""" + + raw_api_response: bytes | None = None + """Raw bytes of the map payload from the device. + + This should be treated as an opaque blob used only internally by this + library to re-parse the map data when needed. + """ + + def __repr__(self) -> str: + img = self.image_content + if img and len(img) > _TRUNCATE_LENGTH: + img = img[: _TRUNCATE_LENGTH - 3] + b"..." + return f"MapContent(image_content={img!r}, map_data={self.map_data!r})" + + +class MapContentTrait(MapContent, Trait): + """Trait for fetching parsed map content for Q7 devices.""" + + def __init__( + self, + map_trait: MapTrait, + *, + serial: str, + model: str, + map_parser_config: B01MapParserConfig | None = None, + ) -> None: + super().__init__() + self._map_trait = map_trait + self._serial = serial + self._model = model + self._map_parser = B01MapParser(map_parser_config) + + async def refresh(self) -> None: + """Fetch, decode, and parse the current map payload.""" + raw_payload = await self._map_trait.get_current_map_payload() + parsed = self.parse_map_content(raw_payload) + self.image_content = parsed.image_content + self.map_data = parsed.map_data + self.raw_api_response = parsed.raw_api_response + + def parse_map_content(self, response: bytes) -> MapContent: + """Parse map content from raw bytes. + + This mirrors the v1 trait behavior so cached map payload bytes can be + reparsed without going back to the device. + """ + try: + parsed_data = self._map_parser.parse( + response, + serial=self._serial, + model=self._model, + ) + except RoborockException: + raise + except Exception as ex: + raise RoborockException("Failed to parse B01 map data") from ex + + if parsed_data.image_content is None: + raise RoborockException("Failed to render B01 map image") + + return MapContent( + image_content=parsed_data.image_content, + map_data=parsed_data.map_data, + raw_api_response=response, + ) diff --git a/roborock/devices/traits/v1/child_lock.py b/roborock/devices/traits/v1/child_lock.py index 11306608..2295026e 100644 --- a/roborock/devices/traits/v1/child_lock.py +++ b/roborock/devices/traits/v1/child_lock.py @@ -9,6 +9,7 @@ class ChildLockTrait(ChildLockStatus, common.V1TraitMixin, common.RoborockSwitch """Trait for controlling the child lock of a Roborock device.""" command = RoborockCommand.GET_CHILD_LOCK_STATUS + converter = common.DefaultConverter(ChildLockStatus) requires_feature = "is_set_child_supported" @property diff --git a/roborock/devices/traits/v1/clean_summary.py b/roborock/devices/traits/v1/clean_summary.py index 6a2fdb7c..61fa4033 100644 --- a/roborock/devices/traits/v1/clean_summary.py +++ b/roborock/devices/traits/v1/clean_summary.py @@ -1,7 +1,6 @@ import logging -from typing import Self -from roborock.data import CleanRecord, CleanSummaryWithDetail +from roborock.data import CleanRecord, CleanSummaryWithDetail, RoborockBase from roborock.devices.traits.v1 import common from roborock.roborock_typing import RoborockCommand from roborock.util import unpack_list @@ -9,48 +8,30 @@ _LOGGER = logging.getLogger(__name__) -class CleanSummaryTrait(CleanSummaryWithDetail, common.V1TraitMixin): - """Trait for managing the clean summary of Roborock devices.""" - - command = RoborockCommand.GET_CLEAN_SUMMARY - - async def refresh(self) -> None: - """Refresh the clean summary data and last clean record. - - Assumes that the clean summary has already been fetched. - """ - await super().refresh() - if not self.records: - _LOGGER.debug("No clean records available in clean summary.") - self.last_clean_record = None - return - last_record_id = self.records[0] - self.last_clean_record = await self.get_clean_record(last_record_id) +class CleanSummaryConverter(common.V1TraitDataConverter): + """Converter for CleanSummaryWithDetail objects.""" - @classmethod - def _parse_type_response(cls, response: common.V1ResponseData) -> Self: + def convert(self, response: common.V1ResponseData) -> RoborockBase: """Parse the response from the device into a CleanSummary.""" if isinstance(response, dict): - return cls.from_dict(response) + return CleanSummaryWithDetail.from_dict(response) elif isinstance(response, list): clean_time, clean_area, clean_count, records = unpack_list(response, 4) - return cls( + return CleanSummaryWithDetail( clean_time=clean_time, clean_area=clean_area, clean_count=clean_count, records=records, ) elif isinstance(response, int): - return cls(clean_time=response) + return CleanSummaryWithDetail(clean_time=response) raise ValueError(f"Unexpected clean summary format: {response!r}") - async def get_clean_record(self, record_id: int) -> CleanRecord: - """Load a specific clean record by ID.""" - response = await self.rpc_channel.send_command(RoborockCommand.GET_CLEAN_RECORD, params=[record_id]) - return self._parse_clean_record_response(response) - @classmethod - def _parse_clean_record_response(cls, response: common.V1ResponseData) -> CleanRecord: +class CleanRecordConverter(common.V1TraitDataConverter): + """Convert server responses to a CleanRecord.""" + + def convert(self, response: common.V1ResponseData) -> CleanRecord: """Parse the response from the device into a CleanRecord.""" if isinstance(response, list) and len(response) == 1: response = response[0] @@ -81,3 +62,29 @@ def _parse_clean_record_response(cls, response: common.V1ResponseData) -> CleanR begin, end, duration, area = unpack_list(response, 4) return CleanRecord(begin=begin, end=end, duration=duration, area=area) raise ValueError(f"Unexpected clean record format: {response!r}") + + +class CleanSummaryTrait(CleanSummaryWithDetail, common.V1TraitMixin): + """Trait for managing the clean summary of Roborock devices.""" + + command = RoborockCommand.GET_CLEAN_SUMMARY + converter = CleanSummaryConverter() + clean_record_converter = CleanRecordConverter() + + async def refresh(self) -> None: + """Refresh the clean summary data and last clean record. + + Assumes that the clean summary has already been fetched. + """ + await super().refresh() + if not self.records: + _LOGGER.debug("No clean records available in clean summary.") + self.last_clean_record = None + return + last_record_id = self.records[0] + self.last_clean_record = await self.get_clean_record(last_record_id) + + async def get_clean_record(self, record_id: int) -> CleanRecord: + """Load a specific clean record by ID.""" + response = await self.rpc_channel.send_command(RoborockCommand.GET_CLEAN_RECORD, params=[record_id]) + return self.clean_record_converter.convert(response) diff --git a/roborock/devices/traits/v1/common.py b/roborock/devices/traits/v1/common.py index 63ae2e20..ce2c899e 100644 --- a/roborock/devices/traits/v1/common.py +++ b/roborock/devices/traits/v1/common.py @@ -5,8 +5,8 @@ import logging from abc import ABC, abstractmethod -from dataclasses import dataclass, fields -from typing import ClassVar, Self +from dataclasses import fields +from typing import ClassVar from roborock.data import RoborockBase from roborock.protocols.v1_protocol import V1RpcChannel @@ -14,10 +14,24 @@ _LOGGER = logging.getLogger(__name__) + V1ResponseData = dict | list | int | str -@dataclass +class V1TraitDataConverter(ABC): + """Converts responses to RoborockBase objects. + + This is an internal class and should not be used directly by consumers. + """ + + @abstractmethod + def convert(self, response: V1ResponseData) -> RoborockBase: + """Convert the values to a dict that can be parsed as a RoborockBase.""" + + def __repr__(self) -> str: + return self.__class__.__name__ + + class V1TraitMixin(ABC): """Base model that supports v1 traits. @@ -42,37 +56,13 @@ class V1TraitMixin(ABC): """ command: ClassVar[RoborockCommand] + """The RoborockCommand used to fetch the trait data from the device (internal only).""" - @classmethod - def _parse_type_response(cls, response: V1ResponseData) -> RoborockBase: - """Parse the response from the device into a a RoborockBase. - - Subclasses should override this method to implement custom parsing - logic as needed. - """ - if not issubclass(cls, RoborockBase): - raise NotImplementedError(f"Trait {cls} does not implement RoborockBase") - # Subclasses can override to implement custom parsing logic - if isinstance(response, list): - response = response[0] - if not isinstance(response, dict): - raise ValueError(f"Unexpected {cls} response format: {response!r}") - return cls.from_dict(response) - - def _parse_response(self, response: V1ResponseData) -> RoborockBase: - """Parse the response from the device into a a RoborockBase. - - This is used by subclasses that want to override the class - behavior with instance-specific data. - """ - return self._parse_type_response(response) - - def __post_init__(self) -> None: - """Post-initialization to set up the RPC channel. + converter: V1TraitDataConverter + """The converter used to parse the response from the device (internal only).""" - This is called automatically after the dataclass is initialized by the - device setup code. - """ + def __init__(self) -> None: + """Initialize the V1TraitMixin.""" self._rpc_channel = None @property @@ -85,32 +75,42 @@ def rpc_channel(self) -> V1RpcChannel: async def refresh(self) -> None: """Refresh the contents of this trait.""" response = await self.rpc_channel.send_command(self.command) - new_data = self._parse_response(response) - if not isinstance(new_data, RoborockBase): - raise ValueError(f"Internal error, unexpected response type: {new_data!r}") - _LOGGER.debug("Refreshed %s: %s", self.__class__.__name__, new_data) - self._update_trait_values(new_data) - - def _update_trait_values(self, new_data: RoborockBase) -> None: - """Update the values of this trait from another instance.""" - for field in fields(new_data): - new_value = getattr(new_data, field.name, None) - setattr(self, field.name, new_value) - - -def _get_value_field(clazz: type[V1TraitMixin]) -> str: - """Get the name of the field marked as the main value of the RoborockValueBase.""" - value_fields = [field.name for field in fields(clazz) if field.metadata.get("roborock_value", False)] - if len(value_fields) != 1: - raise ValueError( - f"RoborockValueBase subclass {clazz} must have exactly one field marked as roborock_value, " - f" but found: {value_fields}" - ) - return value_fields[0] - - -@dataclass(init=False, kw_only=True) -class RoborockValueBase(V1TraitMixin, RoborockBase): + new_data = self.converter.convert(response) + merge_trait_values(self, new_data) # type: ignore[arg-type] + + +def merge_trait_values(target: RoborockBase, new_object: RoborockBase) -> bool: + """Update the target object with set fields in new_object.""" + updated = False + for field in fields(new_object): + old_value = getattr(target, field.name, None) + new_value = getattr(new_object, field.name, None) + if new_value != old_value: + setattr(target, field.name, new_value) + updated = True + return updated + + +class DefaultConverter(V1TraitDataConverter): + """Converts responses to RoborockBase objects.""" + + def __init__(self, dataclass_type: type[RoborockBase]) -> None: + """Initialize the converter.""" + self._dataclass_type = dataclass_type + + def convert(self, response: V1ResponseData) -> RoborockBase: + """Convert the values to a dict that can be parsed as a RoborockBase. + + Subclasses can override to implement custom parsing logic + """ + if isinstance(response, list): + response = response[0] + if not isinstance(response, dict): + raise ValueError(f"Unexpected {self._dataclass_type.__name__} response format: {response!r}") + return self._dataclass_type.from_dict(response) + + +class SingleValueConverter(DefaultConverter): """Base class for traits that represent a single value. This class is intended to be subclassed by traits that represent a single @@ -119,15 +119,18 @@ class RoborockValueBase(V1TraitMixin, RoborockBase): represents the main value of the trait. """ - @classmethod - def _parse_response(cls, response: V1ResponseData) -> Self: + def __init__(self, dataclass_type: type[RoborockBase], value_field: str) -> None: + """Initialize the converter.""" + super().__init__(dataclass_type) + self._value_field = value_field + + def convert(self, response: V1ResponseData) -> RoborockBase: """Parse the response from the device into a RoborockValueBase.""" if isinstance(response, list): response = response[0] if not isinstance(response, int): raise ValueError(f"Unexpected response format: {response!r}") - value_field = _get_value_field(cls) - return cls(**{value_field: response}) + return super().convert({self._value_field: response}) class RoborockSwitchBase(ABC): diff --git a/roborock/devices/traits/v1/consumeable.py b/roborock/devices/traits/v1/consumeable.py index 262f47ca..0e59432f 100644 --- a/roborock/devices/traits/v1/consumeable.py +++ b/roborock/devices/traits/v1/consumeable.py @@ -41,6 +41,7 @@ class ConsumableTrait(Consumable, common.V1TraitMixin): """ command = RoborockCommand.GET_CONSUMABLE + converter = common.DefaultConverter(Consumable) async def reset_consumable(self, consumable: ConsumableAttribute) -> None: """Reset a specific consumable attribute on the device.""" diff --git a/roborock/devices/traits/v1/device_features.py b/roborock/devices/traits/v1/device_features.py index bc619b00..ed6e186e 100644 --- a/roborock/devices/traits/v1/device_features.py +++ b/roborock/devices/traits/v1/device_features.py @@ -8,15 +8,37 @@ from roborock.roborock_typing import RoborockCommand +class DeviceTraitsConverter(common.V1TraitDataConverter): + """Converter for APP_GET_INIT_STATUS responses into DeviceFeatures.""" + + def __init__(self, product: HomeDataProduct) -> None: + """Initialize DeviceTraitsConverter.""" + self._product = product + + def convert(self, response: common.V1ResponseData) -> DeviceFeatures: + """Parse an APP_GET_INIT_STATUS response into a DeviceFeatures instance.""" + if not isinstance(response, list): + raise ValueError(f"Unexpected AppInitStatus response format: {type(response)}: {response!r}") + app_status = AppInitStatus.from_dict(response[0]) + return DeviceFeatures.from_feature_flags( + new_feature_info=app_status.new_feature_info, + new_feature_info_str=app_status.new_feature_info_str, + feature_info=app_status.feature_info, + product_nickname=self._product.product_nickname, + ) + + class DeviceFeaturesTrait(DeviceFeatures, common.V1TraitMixin): """Trait for managing supported features on Roborock devices.""" command = RoborockCommand.APP_GET_INIT_STATUS + converter: DeviceTraitsConverter def __init__(self, product: HomeDataProduct, device_cache: DeviceCache) -> None: # pylint: disable=super-init-not-called """Initialize DeviceFeaturesTrait.""" + common.V1TraitMixin.__init__(self) + self.converter = DeviceTraitsConverter(product) self._product = product - self._nickname = product.product_nickname self._device_cache = device_cache # All fields of DeviceFeatures are required. Initialize them to False # so we have some known state. @@ -54,21 +76,9 @@ async def refresh(self) -> None: """ cache_data = await self._device_cache.get() if cache_data.device_features is not None: - self._update_trait_values(cache_data.device_features) + common.merge_trait_values(self, cache_data.device_features) return # Save cached device features await super().refresh() cache_data.device_features = self await self._device_cache.set(cache_data) - - def _parse_response(self, response: common.V1ResponseData) -> DeviceFeatures: - """Parse the response from the device into a MapContentTrait instance.""" - if not isinstance(response, list): - raise ValueError(f"Unexpected AppInitStatus response format: {type(response)}") - app_status = AppInitStatus.from_dict(response[0]) - return DeviceFeatures.from_feature_flags( - new_feature_info=app_status.new_feature_info, - new_feature_info_str=app_status.new_feature_info_str, - feature_info=app_status.feature_info, - product_nickname=self._nickname, - ) diff --git a/roborock/devices/traits/v1/do_not_disturb.py b/roborock/devices/traits/v1/do_not_disturb.py index 43de8ea5..4865484f 100644 --- a/roborock/devices/traits/v1/do_not_disturb.py +++ b/roborock/devices/traits/v1/do_not_disturb.py @@ -9,6 +9,7 @@ class DoNotDisturbTrait(DnDTimer, common.V1TraitMixin, common.RoborockSwitchBase """Trait for managing Do Not Disturb (DND) settings on Roborock devices.""" command = RoborockCommand.GET_DND_TIMER + converter = common.DefaultConverter(DnDTimer) @property def is_on(self) -> bool: diff --git a/roborock/devices/traits/v1/dust_collection_mode.py b/roborock/devices/traits/v1/dust_collection_mode.py index 00337fff..afda53c1 100644 --- a/roborock/devices/traits/v1/dust_collection_mode.py +++ b/roborock/devices/traits/v1/dust_collection_mode.py @@ -10,4 +10,5 @@ class DustCollectionModeTrait(DustCollectionMode, common.V1TraitMixin): """Trait for dust collection mode.""" command = RoborockCommand.GET_DUST_COLLECTION_MODE + converter = common.DefaultConverter(DustCollectionMode) requires_dock_type = is_valid_dock diff --git a/roborock/devices/traits/v1/flow_led_status.py b/roborock/devices/traits/v1/flow_led_status.py index 5a406000..3c1ae15d 100644 --- a/roborock/devices/traits/v1/flow_led_status.py +++ b/roborock/devices/traits/v1/flow_led_status.py @@ -9,6 +9,7 @@ class FlowLedStatusTrait(FlowLedStatus, common.V1TraitMixin, common.RoborockSwit """Trait for controlling the Flow LED status of a Roborock device.""" command = RoborockCommand.GET_FLOW_LED_STATUS + converter = common.DefaultConverter(FlowLedStatus) requires_feature = "is_flow_led_setting_supported" @property diff --git a/roborock/devices/traits/v1/home.py b/roborock/devices/traits/v1/home.py index cbe1fc1e..5034d33c 100644 --- a/roborock/devices/traits/v1/home.py +++ b/roborock/devices/traits/v1/home.py @@ -18,7 +18,6 @@ import asyncio import base64 import logging -from typing import Self from roborock.data import CombinedMapInfo, MultiMapsListMapInfo, NamedRoomMapping, RoborockBase from roborock.data.v1.v1_code_mappings import RoborockStateCode @@ -41,6 +40,7 @@ class HomeTrait(RoborockBase, common.V1TraitMixin): """Trait that represents a full view of the home layout.""" command = RoborockCommand.GET_MAP_V1 # This is not used + converter = common.DefaultConverter(RoborockBase) # Not used def __init__( self, @@ -93,7 +93,7 @@ async def discover_home(self) -> None: self._discovery_completed = True try: self._home_map_content = { - k: self._map_content.parse_map_content(base64.b64decode(v)) + k: self._map_content.converter.parse_map_content(base64.b64decode(v)) for k, v in (device_cache_data.home_map_content_base64 or {}).items() } except (ValueError, RoborockException) as ex: @@ -233,10 +233,6 @@ def home_map_content(self) -> dict[int, MapContent] | None: """Returns the map content for all cached maps.""" return self._home_map_content - def _parse_response(self, response: common.V1ResponseData) -> Self: - """This trait does not parse responses directly.""" - raise NotImplementedError("HomeTrait does not support direct command responses") - async def _update_home_cache( self, home_map_info: dict[int, CombinedMapInfo], home_map_content: dict[int, MapContent] ) -> None: diff --git a/roborock/devices/traits/v1/led_status.py b/roborock/devices/traits/v1/led_status.py index 41d15d53..98cd1cbf 100644 --- a/roborock/devices/traits/v1/led_status.py +++ b/roborock/devices/traits/v1/led_status.py @@ -5,10 +5,28 @@ from .common import V1ResponseData +class LedStatusConverter(common.V1TraitDataConverter): + """Converter for LedStatus.""" + + def convert(self, response: V1ResponseData) -> LedStatus: + """Parse the response from the device into a a RoborockBase. + + Subclasses should override this method to implement custom parsing + logic as needed. + """ + if not isinstance(response, list): + raise ValueError(f"Unexpected LedStatus response format: {response!r}") + response = response[0] + if not isinstance(response, int): + raise ValueError(f"Unexpected LedStatus response format: {response!r}") + return LedStatus.from_dict({"status": response}) + + class LedStatusTrait(LedStatus, common.V1TraitMixin, common.RoborockSwitchBase): """Trait for controlling the LED status of a Roborock device.""" command = RoborockCommand.GET_LED_STATUS + converter = LedStatusConverter() requires_feature = "is_led_status_switch_supported" @property @@ -27,17 +45,3 @@ async def disable(self) -> None: await self.rpc_channel.send_command(RoborockCommand.SET_LED_STATUS, params=[0]) # Optimistic update to avoid an extra refresh self.status = 0 - - @classmethod - def _parse_type_response(cls, response: V1ResponseData) -> LedStatus: - """Parse the response from the device into a a RoborockBase. - - Subclasses should override this method to implement custom parsing - logic as needed. - """ - if not isinstance(response, list): - raise ValueError(f"Unexpected {cls} response format: {response!r}") - response = response[0] - if not isinstance(response, int): - raise ValueError(f"Unexpected {cls} response format: {response!r}") - return cls.from_dict({"status": response}) diff --git a/roborock/devices/traits/v1/map_content.py b/roborock/devices/traits/v1/map_content.py index e0cc3e47..35b015d6 100644 --- a/roborock/devices/traits/v1/map_content.py +++ b/roborock/devices/traits/v1/map_content.py @@ -40,19 +40,15 @@ def __repr__(self) -> str: return f"MapContent(image_content={img!r}, map_data={self.map_data!r})" -@common.map_rpc_channel -class MapContentTrait(MapContent, common.V1TraitMixin): - """Trait for fetching the map content.""" - - command = RoborockCommand.GET_MAP_V1 +class MapContentConverter(common.V1TraitDataConverter): + """Convert map response data to MapContent.""" - def __init__(self, map_parser_config: MapParserConfig | None = None) -> None: - """Initialize MapContentTrait.""" - super().__init__() - self._map_parser = MapParser(map_parser_config or MapParserConfig()) + def __init__(self, map_parser: MapParser) -> None: + """Initialize MapContentConverter.""" + self._map_parser = map_parser - def _parse_response(self, response: common.V1ResponseData) -> MapContent: - """Parse the response from the device into a MapContentTrait instance.""" + def convert(self, response: common.V1ResponseData) -> MapContent: + """Parse the response from the device into a MapContent instance.""" if not isinstance(response, bytes): raise ValueError(f"Unexpected MapContentTrait response format: {type(response)}") return self.parse_map_content(response) @@ -81,3 +77,16 @@ def parse_map_content(self, response: bytes) -> MapContent: map_data=parsed_data.map_data, raw_api_response=response, ) + + +@common.map_rpc_channel +class MapContentTrait(MapContent, common.V1TraitMixin): + """Trait for fetching the map content.""" + + command = RoborockCommand.GET_MAP_V1 + converter: MapContentConverter + + def __init__(self, map_parser_config: MapParserConfig | None = None) -> None: + """Initialize MapContentTrait.""" + super().__init__() + self.converter = MapContentConverter(MapParser(map_parser_config or MapParserConfig())) diff --git a/roborock/devices/traits/v1/maps.py b/roborock/devices/traits/v1/maps.py index 0f20b58a..4349bd8d 100644 --- a/roborock/devices/traits/v1/maps.py +++ b/roborock/devices/traits/v1/maps.py @@ -6,7 +6,6 @@ """ import logging -from typing import Self from roborock.data import MultiMapsList, MultiMapsListMapInfo from roborock.devices.traits.v1 import common @@ -17,6 +16,24 @@ _LOGGER = logging.getLogger(__name__) +class MultiMapsListConverter(common.V1TraitDataConverter): + """Converters responses to MultiMapsList.""" + + def convert(self, response: common.V1ResponseData) -> MultiMapsList: + """Parse the response from the device into a MapsTrait instance. + + This overrides the base implementation to handle the specific + response format for the multi maps list. This is needed because we have + a custom constructor that requires the StatusTrait. + """ + if not isinstance(response, list): + raise ValueError(f"Unexpected MapsTrait response format: {response!r}") + response = response[0] + if not isinstance(response, dict): + raise ValueError(f"Unexpected MapsTrait response format: {response!r}") + return MultiMapsList.from_dict(response) + + @common.mqtt_rpc_channel class MapsTrait(MultiMapsList, common.V1TraitMixin): """Trait for managing the maps of Roborock devices. @@ -34,6 +51,7 @@ class MapsTrait(MultiMapsList, common.V1TraitMixin): """ command = RoborockCommand.GET_MULTI_MAPS_LIST + converter = MultiMapsListConverter() def __init__(self, status_trait: StatusTrait) -> None: """Initialize the MapsTrait. @@ -64,17 +82,3 @@ async def set_current_map(self, map_flag: int) -> None: await self.rpc_channel.send_command(RoborockCommand.LOAD_MULTI_MAP, params=[map_flag]) # Refresh our status to make sure it reflects the new map await self._status_trait.refresh() - - def _parse_response(self, response: common.V1ResponseData) -> Self: - """Parse the response from the device into a MapsTrait instance. - - This overrides the base implementation to handle the specific - response format for the multi maps list. This is needed because we have - a custom constructor that requires the StatusTrait. - """ - if not isinstance(response, list): - raise ValueError(f"Unexpected MapsTrait response format: {response!r}") - response = response[0] - if not isinstance(response, dict): - raise ValueError(f"Unexpected MapsTrait response format: {response!r}") - return MultiMapsList.from_dict(response) diff --git a/roborock/devices/traits/v1/network_info.py b/roborock/devices/traits/v1/network_info.py index a88394cd..2bdca57a 100644 --- a/roborock/devices/traits/v1/network_info.py +++ b/roborock/devices/traits/v1/network_info.py @@ -1,7 +1,5 @@ """Trait for device network information.""" -from __future__ import annotations - import logging from roborock.data import NetworkInfo @@ -12,6 +10,16 @@ _LOGGER = logging.getLogger(__name__) +class NetworkInfoConverter(common.V1TraitDataConverter): + """Converter for NetworkInfo objects.""" + + def convert(self, response: common.V1ResponseData) -> NetworkInfo: + """Parse the response from the device into a NetworkInfoConverter instance.""" + if not isinstance(response, dict): + raise ValueError(f"Unexpected NetworkInfoTrait response format: {response!r}") + return NetworkInfo.from_dict(response) + + class NetworkInfoTrait(NetworkInfo, common.V1TraitMixin): """Trait for device network information. @@ -23,6 +31,7 @@ class NetworkInfoTrait(NetworkInfo, common.V1TraitMixin): """ command = RoborockCommand.GET_NETWORK_INFO + converter = NetworkInfoConverter() def __init__(self, device_uid: str, device_cache: DeviceCache) -> None: # pylint: disable=super-init-not-called """Initialize the trait.""" @@ -36,7 +45,7 @@ async def refresh(self) -> None: device_cache_data = await self._device_cache.get() if device_cache_data.network_info: _LOGGER.debug("Using cached network info for device %s", self._device_uid) - self._update_trait_values(device_cache_data.network_info) + common.merge_trait_values(self, device_cache_data.network_info) return # Load from device if not in cache @@ -47,9 +56,3 @@ async def refresh(self) -> None: device_cache_data = await self._device_cache.get() device_cache_data.network_info = self await self._device_cache.set(device_cache_data) - - def _parse_response(self, response: common.V1ResponseData) -> NetworkInfo: - """Parse the response from the device into a NetworkInfo.""" - if not isinstance(response, dict): - raise ValueError(f"Unexpected NetworkInfoTrait response format: {response!r}") - return NetworkInfo.from_dict(response) diff --git a/roborock/devices/traits/v1/rooms.py b/roborock/devices/traits/v1/rooms.py index 8d90c41c..6c6e35c7 100644 --- a/roborock/devices/traits/v1/rooms.py +++ b/roborock/devices/traits/v1/rooms.py @@ -25,11 +25,63 @@ def room_map(self) -> dict[int, NamedRoomMapping]: return {} return {room.segment_id: room for room in self.rooms} + def with_room_names(self, name_map: dict[str, str]) -> "Rooms": + """Create a new Rooms object with updated room names.""" + return Rooms( + rooms=[ + NamedRoomMapping( + segment_id=room.segment_id, + iot_id=room.iot_id, + raw_name=name_map.get(room.iot_id), + ) + for room in self.rooms or [] + ] + ) + + +class RoomsConverter(common.V1TraitDataConverter): + """Converts response objects to Rooms.""" + + def convert(self, response: common.V1ResponseData) -> Rooms: + """Parse the response from the device into a list of NamedRoomMapping.""" + if not isinstance(response, list): + raise ValueError(f"Unexpected RoomsTrait response format: {response!r}") + segment_map = self.extract_segment_map(response) + return Rooms( + rooms=[NamedRoomMapping(segment_id=segment_id, iot_id=iot_id) for segment_id, iot_id in segment_map.items()] + ) + + @staticmethod + def extract_segment_map(response: list) -> dict[int, str]: + """Extract a segment_id -> iot_id mapping from the response. + + The response format can be either a flat list of [segment_id, iot_id] or a + list of lists, where each inner list is a pair of [segment_id, iot_id]. This + function normalizes the response into a dict of segment_id to iot_id. + + NOTE: We currently only partial samples of the room mapping formats, so + improving test coverage with samples from a real device with this format + would be helpful. + """ + if len(response) == 2 and not isinstance(response[0], list): + segment_id, iot_id = response[0], response[1] + return {segment_id: str(iot_id)} + + segment_map: dict[int, str] = {} + for part in response: + if not isinstance(part, list) or len(part) < 2: + _LOGGER.warning("Unexpected room mapping entry format: %r", part) + continue + segment_id, iot_id = part[0], part[1] + segment_map[segment_id] = str(iot_id) + return segment_map + class RoomsTrait(Rooms, common.V1TraitMixin): """Trait for managing the room mappings of Roborock devices.""" command = RoborockCommand.GET_ROOM_MAPPING + converter = RoomsConverter() def __init__(self, home_data: HomeData, web_api: UserWebApiClient) -> None: """Initialize the RoomsTrait.""" @@ -44,7 +96,7 @@ async def refresh(self) -> None: if not isinstance(response, list): raise ValueError(f"Unexpected RoomsTrait response format: {response!r}") - segment_map = _extract_segment_map(response) + segment_map = RoomsConverter.extract_segment_map(response) # Track all iot ids seen before. Refresh the room list when new ids are found. new_iot_ids = set(segment_map.values()) - set(self._home_data.rooms_map.keys()) if new_iot_ids - self._discovered_iot_ids: @@ -54,22 +106,9 @@ async def refresh(self) -> None: self._home_data.rooms = updated_rooms self._discovered_iot_ids.update(new_iot_ids) - new_data = self._parse_rooms(segment_map, self._home_data.rooms_name_map) - self._update_trait_values(new_data) - _LOGGER.debug("Refreshed %s: %s", self.__class__.__name__, new_data) - - @staticmethod - def _parse_rooms( - segment_map: dict[int, str], - name_map: dict[str, str], - ) -> Rooms: - """Parse the response from the device into a list of NamedRoomMapping.""" - return Rooms( - rooms=[ - NamedRoomMapping(segment_id=segment_id, iot_id=iot_id, raw_name=name_map.get(iot_id)) - for segment_id, iot_id in segment_map.items() - ] - ) + rooms = self.converter.convert(response) + rooms = rooms.with_room_names(self._home_data.rooms_name_map) + common.merge_trait_values(self, rooms) async def _refresh_rooms(self) -> list[HomeDataRoom]: """Fetch the latest rooms from the web API.""" @@ -78,28 +117,3 @@ async def _refresh_rooms(self) -> list[HomeDataRoom]: except Exception: _LOGGER.debug("Failed to fetch rooms from web API", exc_info=True) return [] - - -def _extract_segment_map(response: list) -> dict[int, str]: - """Extract a segment_id -> iot_id mapping from the response. - - The response format can be either a flat list of [segment_id, iot_id] or a - list of lists, where each inner list is a pair of [segment_id, iot_id]. This - function normalizes the response into a dict of segment_id to iot_id. - - NOTE: We currently only partial samples of the room mapping formats, so - improving test coverage with samples from a real device with this format - would be helpful. - """ - if len(response) == 2 and not isinstance(response[0], list): - segment_id, iot_id = response[0], response[1] - return {segment_id: str(iot_id)} - - segment_map: dict[int, str] = {} - for part in response: - if not isinstance(part, list) or len(part) < 2: - _LOGGER.warning("Unexpected room mapping entry format: %r", part) - continue - segment_id, iot_id = part[0], part[1] - segment_map[segment_id] = str(iot_id) - return segment_map diff --git a/roborock/devices/traits/v1/smart_wash_params.py b/roborock/devices/traits/v1/smart_wash_params.py index cb71c070..9df0a53d 100644 --- a/roborock/devices/traits/v1/smart_wash_params.py +++ b/roborock/devices/traits/v1/smart_wash_params.py @@ -10,4 +10,5 @@ class SmartWashParamsTrait(SmartWashParams, common.V1TraitMixin): """Trait for smart wash parameters.""" command = RoborockCommand.GET_SMART_WASH_PARAMS + converter = common.DefaultConverter(SmartWashParams) requires_dock_type = is_wash_n_fill_dock diff --git a/roborock/devices/traits/v1/status.py b/roborock/devices/traits/v1/status.py index 961d2110..82371c15 100644 --- a/roborock/devices/traits/v1/status.py +++ b/roborock/devices/traits/v1/status.py @@ -1,5 +1,4 @@ from functools import cached_property -from typing import Self from roborock import ( CleanRoutes, @@ -43,6 +42,7 @@ class StatusTrait(StatusV2, common.V1TraitMixin): """ command = RoborockCommand.GET_STATUS + converter = common.DefaultConverter(StatusV2) def __init__(self, device_feature_trait: DeviceFeaturesTrait, region: str | None = None) -> None: """Initialize the StatusTrait.""" @@ -91,11 +91,3 @@ def mop_route_name(self) -> str | None: if self.mop_mode is None: return None return self.mop_route_mapping.get(self.mop_mode) - - def _parse_response(self, response: common.V1ResponseData) -> Self: - """Parse the response from the device into a StatusV2-based status object.""" - if isinstance(response, list): - response = response[0] - if isinstance(response, dict): - return StatusV2.from_dict(response) - raise ValueError(f"Unexpected status format: {response!r}") diff --git a/roborock/devices/traits/v1/valley_electricity_timer.py b/roborock/devices/traits/v1/valley_electricity_timer.py index d21a229d..07623d4c 100644 --- a/roborock/devices/traits/v1/valley_electricity_timer.py +++ b/roborock/devices/traits/v1/valley_electricity_timer.py @@ -9,6 +9,7 @@ class ValleyElectricityTimerTrait(ValleyElectricityTimer, common.V1TraitMixin, c """Trait for managing Valley Electricity Timer settings on Roborock devices.""" command = RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER + converter = common.DefaultConverter(ValleyElectricityTimer) requires_feature = "is_supported_valley_electricity" @property diff --git a/roborock/devices/traits/v1/volume.py b/roborock/devices/traits/v1/volume.py index 865a4b3b..0655b821 100644 --- a/roborock/devices/traits/v1/volume.py +++ b/roborock/devices/traits/v1/volume.py @@ -1,18 +1,15 @@ -from dataclasses import dataclass, field +from dataclasses import dataclass +from roborock.data.containers import RoborockBase from roborock.devices.traits.v1 import common from roborock.roborock_typing import RoborockCommand -# TODO: This is currently the pattern for holding all the commands that hold a -# single value, but it still seems too verbose. Maybe we can generate these -# dynamically or somehow make them less code. - @dataclass -class SoundVolume(common.RoborockValueBase): +class SoundVolume(RoborockBase): """Dataclass for sound volume.""" - volume: int | None = field(default=None, metadata={"roborock_value": True}) + volume: int | None = None """Sound volume level (0-100).""" @@ -20,6 +17,7 @@ class SoundVolumeTrait(SoundVolume, common.V1TraitMixin): """Trait for controlling the sound volume of a Roborock device.""" command = RoborockCommand.GET_SOUND_VOLUME + converter = common.SingleValueConverter(SoundVolume, "volume") async def set_volume(self, volume: int) -> None: """Set the sound volume of the device.""" diff --git a/roborock/devices/traits/v1/wash_towel_mode.py b/roborock/devices/traits/v1/wash_towel_mode.py index 9dcbf63f..de607bc7 100644 --- a/roborock/devices/traits/v1/wash_towel_mode.py +++ b/roborock/devices/traits/v1/wash_towel_mode.py @@ -14,6 +14,7 @@ class WashTowelModeTrait(WashTowelMode, common.V1TraitMixin): """Trait for wash towel mode.""" command = RoborockCommand.GET_WASH_TOWEL_MODE + converter = common.DefaultConverter(WashTowelMode) requires_dock_type = is_wash_n_fill_dock def __init__( diff --git a/roborock/diagnostics.py b/roborock/diagnostics.py index f455f607..c238209c 100644 --- a/roborock/diagnostics.py +++ b/roborock/diagnostics.py @@ -9,13 +9,11 @@ DeviceManager. """ -from __future__ import annotations - import time from collections import Counter from collections.abc import Generator, Mapping from contextlib import contextmanager -from typing import Any, TypeVar, cast +from typing import Any, Self, TypeVar, cast class Diagnostics: @@ -28,7 +26,7 @@ class Diagnostics: def __init__(self) -> None: """Initialize Diagnostics.""" self._counter: Counter = Counter() - self._subkeys: dict[str, Diagnostics] = {} + self._subkeys: dict[str, Self] = {} def increment(self, key: str, count: int = 1) -> None: """Increment a counter for the specified key/event.""" @@ -49,7 +47,7 @@ def as_dict(self) -> Mapping[str, Any]: data[k] = v return data - def subkey(self, key: str) -> Diagnostics: + def subkey(self, key: str) -> Self: """Return sub-Diagnostics object with the specified subkey. This will create a new Diagnostics object if one does not already exist @@ -63,7 +61,7 @@ def subkey(self, key: str) -> Diagnostics: The Diagnostics object for the specified subkey. """ if key not in self._subkeys: - self._subkeys[key] = Diagnostics() + self._subkeys[key] = type(self)() return self._subkeys[key] @contextmanager diff --git a/roborock/exceptions.py b/roborock/exceptions.py index 3c5b8295..55e72b60 100644 --- a/roborock/exceptions.py +++ b/roborock/exceptions.py @@ -1,7 +1,5 @@ """Roborock exceptions.""" -from __future__ import annotations - class RoborockException(Exception): """Class for Roborock exceptions.""" diff --git a/roborock/map/b01_map_parser.py b/roborock/map/b01_map_parser.py new file mode 100644 index 00000000..57b5534a --- /dev/null +++ b/roborock/map/b01_map_parser.py @@ -0,0 +1,188 @@ +"""Module for parsing B01/Q7 map content. + +Observed Q7 `MAP_RESPONSE` payloads follow this decode pipeline: +- base64-encoded ASCII +- AES-ECB encrypted with the derived map key +- PKCS7 padded +- ASCII hex for a zlib-compressed SCMap payload + +The inner SCMap blob is parsed with protobuf messages generated from +`roborock/map/proto/b01_scmap.proto`. +""" + +import base64 +import binascii +import hashlib +import io +import zlib +from dataclasses import dataclass + +from Crypto.Cipher import AES +from Crypto.Util.Padding import pad, unpad +from google.protobuf.message import DecodeError, Message +from PIL import Image +from vacuum_map_parser_base.config.image_config import ImageConfig +from vacuum_map_parser_base.map_data import ImageData, MapData + +from roborock.exceptions import RoborockException +from roborock.map.proto.b01_scmap_pb2 import RobotMap # type: ignore[attr-defined] + +from .map_parser import ParsedMapData + +_B64_CHARS = set(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=") +_MAP_FILE_FORMAT = "PNG" + + +@dataclass +class B01MapParserConfig: + """Configuration for the B01/Q7 map parser.""" + + map_scale: int = 4 + """Scale factor for the rendered map image.""" + + +class B01MapParser: + """Decoder/parser for B01/Q7 SCMap payloads.""" + + def __init__(self, config: B01MapParserConfig | None = None) -> None: + self._config = config or B01MapParserConfig() + + def parse(self, raw_payload: bytes, *, serial: str, model: str) -> ParsedMapData: + """Parse a raw MAP_RESPONSE payload and return a PNG + MapData.""" + inflated = _decode_b01_map_payload(raw_payload, serial=serial, model=model) + parsed = _parse_scmap_payload(inflated) + size_x, size_y, grid = _extract_grid(parsed) + room_names = _extract_room_names(parsed) + + image = _render_occupancy_image(grid, size_x=size_x, size_y=size_y, scale=self._config.map_scale) + + map_data = MapData() + map_data.image = ImageData( + size=size_x * size_y, + top=0, + left=0, + height=size_y, + width=size_x, + image_config=ImageConfig(scale=self._config.map_scale), + data=image, + img_transformation=lambda p: p, + ) + if room_names: + map_data.additional_parameters["room_names"] = room_names + + image_bytes = io.BytesIO() + image.save(image_bytes, format=_MAP_FILE_FORMAT) + + return ParsedMapData( + image_content=image_bytes.getvalue(), + map_data=map_data, + ) + + +def _derive_map_key(serial: str, model: str) -> bytes: + """Derive the B01/Q7 map decrypt key from serial + model.""" + model_suffix = model.split(".")[-1] + model_key = (model_suffix + "0" * 16)[:16].encode() + material = f"{serial}+{model_suffix}+{serial}".encode() + encrypted = AES.new(model_key, AES.MODE_ECB).encrypt(pad(material, AES.block_size)) + md5 = hashlib.md5(base64.b64encode(encrypted), usedforsecurity=False).hexdigest() + return md5[8:24].encode() + + +def _decode_base64_payload(raw_payload: bytes) -> bytes: + blob = raw_payload.strip() + if len(blob) < 32 or any(b not in _B64_CHARS for b in blob): + raise RoborockException("Failed to decode B01 map payload") + + padded = blob + b"=" * (-len(blob) % 4) + try: + return base64.b64decode(padded, validate=True) + except binascii.Error as err: + raise RoborockException("Failed to decode B01 map payload") from err + + +def _decode_b01_map_payload(raw_payload: bytes, *, serial: str, model: str) -> bytes: + """Decode raw B01 `MAP_RESPONSE` payload into inflated SCMap bytes.""" + encrypted_payload = _decode_base64_payload(raw_payload) + if len(encrypted_payload) % AES.block_size != 0: + raise RoborockException("Unexpected encrypted B01 map payload length") + + map_key = _derive_map_key(serial, model) + decrypted_hex = AES.new(map_key, AES.MODE_ECB).decrypt(encrypted_payload) + + try: + compressed_hex = unpad(decrypted_hex, AES.block_size).decode("ascii") + compressed_payload = bytes.fromhex(compressed_hex) + return zlib.decompress(compressed_payload) + except (ValueError, UnicodeDecodeError, zlib.error) as err: + raise RoborockException("Failed to decode B01 map payload") from err + + +def _parse_proto(blob: bytes, message: Message, *, context: str) -> None: + try: + message.ParseFromString(blob) + except DecodeError as err: + raise RoborockException(f"Failed to parse {context}") from err + + +def _decode_map_data_bytes(value: bytes) -> bytes: + try: + return zlib.decompress(value) + except zlib.error: + return value + + +def _parse_scmap_payload(payload: bytes) -> RobotMap: + """Parse inflated SCMap bytes into a generated protobuf message.""" + parsed = RobotMap() + _parse_proto(payload, parsed, context="B01 SCMap") + return parsed + + +def _extract_grid(parsed: RobotMap) -> tuple[int, int, bytes]: + if not parsed.HasField("mapHead") or not parsed.HasField("mapData"): + raise RoborockException("Failed to parse B01 map header/grid") + + size_x = parsed.mapHead.sizeX if parsed.mapHead.HasField("sizeX") else 0 + size_y = parsed.mapHead.sizeY if parsed.mapHead.HasField("sizeY") else 0 + if not size_x or not size_y or not parsed.mapData.HasField("mapData"): + raise RoborockException("Failed to parse B01 map header/grid") + + map_data = _decode_map_data_bytes(parsed.mapData.mapData) + expected_len = size_x * size_y + if len(map_data) < expected_len: + raise RoborockException("B01 map data shorter than expected dimensions") + + return size_x, size_y, map_data[:expected_len] + + +def _extract_room_names(parsed: RobotMap) -> dict[int, str]: + # Expose room id/name mapping without inventing room geometry/polygons. + room_names: dict[int, str] = {} + for room in parsed.roomDataInfo: + if room.HasField("roomId"): + room_id = room.roomId + room_names[room_id] = room.roomName if room.HasField("roomName") else f"Room {room_id}" + return room_names + + +def _render_occupancy_image(grid: bytes, *, size_x: int, size_y: int, scale: int) -> Image.Image: + """Render the B01 occupancy grid into a simple image.""" + + # The observed occupancy grid contains only: + # - 0: outside/unknown + # - 127: wall/obstacle + # - 128: floor/free + table = bytearray(range(256)) + table[0] = 0 + table[127] = 180 + table[128] = 255 + + mapped = grid.translate(bytes(table)) + img = Image.frombytes("L", (size_x, size_y), mapped) + img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM).convert("RGB") + + if scale > 1: + img = img.resize((size_x * scale, size_y * scale), resample=Image.Resampling.NEAREST) + + return img diff --git a/roborock/map/proto/__init__.py b/roborock/map/proto/__init__.py new file mode 100644 index 00000000..81a4f8b1 --- /dev/null +++ b/roborock/map/proto/__init__.py @@ -0,0 +1 @@ +"""Generated protobuf modules for Roborock map payloads.""" diff --git a/roborock/map/proto/b01_scmap.proto b/roborock/map/proto/b01_scmap.proto new file mode 100644 index 00000000..b3659813 --- /dev/null +++ b/roborock/map/proto/b01_scmap.proto @@ -0,0 +1,70 @@ +// Checked-in B01/Q7 SCMap schema for the generated runtime protobuf module. +// Regenerate the checked-in Python module after edits with: +// python -m grpc_tools.protoc -I./roborock/map/proto --python_out=./roborock/map/proto roborock/map/proto/b01_scmap.proto +// The generated file `b01_scmap_pb2.py` is checked in for runtime use and should +// not be edited by hand. +syntax = "proto2"; + +package b01.scmap; + +message DevicePointInfo { + optional float x = 1; + optional float y = 2; +} + +message MapBoundaryInfo { + optional string mapMd5 = 1; + optional uint32 vMinX = 2; + optional uint32 vMaxX = 3; + optional uint32 vMinY = 4; + optional uint32 vMaxY = 5; +} + +message MapExtInfo { + optional uint32 taskBeginDate = 1; + optional uint32 mapUploadDate = 2; + optional uint32 mapValid = 3; + optional uint32 radian = 4; + optional uint32 force = 5; + optional uint32 cleanPath = 6; + optional MapBoundaryInfo boudaryInfo = 7; + optional uint32 mapVersion = 8; + optional uint32 mapValueType = 9; +} + +message MapHeadInfo { + optional uint32 mapHeadId = 1; + optional uint32 sizeX = 2; + optional uint32 sizeY = 3; + optional float minX = 4; + optional float minY = 5; + optional float maxX = 6; + optional float maxY = 7; + optional float resolution = 8; +} + +message MapDataInfo { + optional bytes mapData = 1; +} + +message RoomDataInfo { + optional uint32 roomId = 1; + optional string roomName = 2; + optional uint32 roomTypeId = 3; + optional uint32 meterialId = 4; + optional uint32 cleanState = 5; + optional uint32 roomClean = 6; + optional uint32 roomCleanIndex = 7; + optional DevicePointInfo roomNamePost = 8; + optional uint32 colorId = 10; + optional uint32 floor_direction = 11; + optional uint32 global_seq = 12; +} + +message RobotMap { + optional uint32 mapType = 1; + optional MapExtInfo mapExtInfo = 2; + optional MapHeadInfo mapHead = 3; + optional MapDataInfo mapData = 4; + repeated RoomDataInfo roomDataInfo = 12; +} diff --git a/roborock/map/proto/b01_scmap_pb2.py b/roborock/map/proto/b01_scmap_pb2.py new file mode 100644 index 00000000..66cc0843 --- /dev/null +++ b/roborock/map/proto/b01_scmap_pb2.py @@ -0,0 +1,48 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: roborock/map/proto/b01_scmap.proto +# Protobuf Python Version: 6.31.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 31, + 1, + '', + 'roborock/map/proto/b01_scmap.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\"roborock/map/proto/b01_scmap.proto\x12\tb01.scmap\"\'\n\x0f\x44\x65vicePointInfo\x12\t\n\x01x\x18\x01 \x01(\x02\x12\t\n\x01y\x18\x02 \x01(\x02\"]\n\x0fMapBoundaryInfo\x12\x0e\n\x06mapMd5\x18\x01 \x01(\t\x12\r\n\x05vMinX\x18\x02 \x01(\r\x12\r\n\x05vMaxX\x18\x03 \x01(\r\x12\r\n\x05vMinY\x18\x04 \x01(\r\x12\r\n\x05vMaxY\x18\x05 \x01(\r\"\xd9\x01\n\nMapExtInfo\x12\x15\n\rtaskBeginDate\x18\x01 \x01(\r\x12\x15\n\rmapUploadDate\x18\x02 \x01(\r\x12\x10\n\x08mapValid\x18\x03 \x01(\r\x12\x0e\n\x06radian\x18\x04 \x01(\r\x12\r\n\x05\x66orce\x18\x05 \x01(\r\x12\x11\n\tcleanPath\x18\x06 \x01(\r\x12/\n\x0b\x62oudaryInfo\x18\x07 \x01(\x0b\x32\x1a.b01.scmap.MapBoundaryInfo\x12\x12\n\nmapVersion\x18\x08 \x01(\r\x12\x14\n\x0cmapValueType\x18\t \x01(\r\"\x8a\x01\n\x0bMapHeadInfo\x12\x11\n\tmapHeadId\x18\x01 \x01(\r\x12\r\n\x05sizeX\x18\x02 \x01(\r\x12\r\n\x05sizeY\x18\x03 \x01(\r\x12\x0c\n\x04minX\x18\x04 \x01(\x02\x12\x0c\n\x04minY\x18\x05 \x01(\x02\x12\x0c\n\x04maxX\x18\x06 \x01(\x02\x12\x0c\n\x04maxY\x18\x07 \x01(\x02\x12\x12\n\nresolution\x18\x08 \x01(\x02\"\x1e\n\x0bMapDataInfo\x12\x0f\n\x07mapData\x18\x01 \x01(\x0c\"\x87\x02\n\x0cRoomDataInfo\x12\x0e\n\x06roomId\x18\x01 \x01(\r\x12\x10\n\x08roomName\x18\x02 \x01(\t\x12\x12\n\nroomTypeId\x18\x03 \x01(\r\x12\x12\n\nmeterialId\x18\x04 \x01(\r\x12\x12\n\ncleanState\x18\x05 \x01(\r\x12\x11\n\troomClean\x18\x06 \x01(\r\x12\x16\n\x0eroomCleanIndex\x18\x07 \x01(\r\x12\x30\n\x0croomNamePost\x18\x08 \x01(\x0b\x32\x1a.b01.scmap.DevicePointInfo\x12\x0f\n\x07\x63olorId\x18\n \x01(\r\x12\x17\n\x0f\x66loor_direction\x18\x0b \x01(\r\x12\x12\n\nglobal_seq\x18\x0c \x01(\r\"\xc7\x01\n\x08RobotMap\x12\x0f\n\x07mapType\x18\x01 \x01(\r\x12)\n\nmapExtInfo\x18\x02 \x01(\x0b\x32\x15.b01.scmap.MapExtInfo\x12\'\n\x07mapHead\x18\x03 \x01(\x0b\x32\x16.b01.scmap.MapHeadInfo\x12\'\n\x07mapData\x18\x04 \x01(\x0b\x32\x16.b01.scmap.MapDataInfo\x12-\n\x0croomDataInfo\x18\x0c \x03(\x0b\x32\x17.b01.scmap.RoomDataInfo') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'roborock.map.proto.b01_scmap_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_DEVICEPOINTINFO']._serialized_start=49 + _globals['_DEVICEPOINTINFO']._serialized_end=88 + _globals['_MAPBOUNDARYINFO']._serialized_start=90 + _globals['_MAPBOUNDARYINFO']._serialized_end=183 + _globals['_MAPEXTINFO']._serialized_start=186 + _globals['_MAPEXTINFO']._serialized_end=403 + _globals['_MAPHEADINFO']._serialized_start=406 + _globals['_MAPHEADINFO']._serialized_end=544 + _globals['_MAPDATAINFO']._serialized_start=546 + _globals['_MAPDATAINFO']._serialized_end=576 + _globals['_ROOMDATAINFO']._serialized_start=579 + _globals['_ROOMDATAINFO']._serialized_end=842 + _globals['_ROBOTMAP']._serialized_start=845 + _globals['_ROBOTMAP']._serialized_end=1044 +# @@protoc_insertion_point(module_scope) diff --git a/roborock/protocol.py b/roborock/protocol.py index 828a432a..fc3af842 100644 --- a/roborock/protocol.py +++ b/roborock/protocol.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import binascii import gzip import hashlib diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py index fbe9becf..7c7f65a6 100644 --- a/roborock/roborock_message.py +++ b/roborock/roborock_message.py @@ -1,7 +1,6 @@ -from __future__ import annotations - from dataclasses import dataclass, field from enum import StrEnum +from typing import Self from roborock import RoborockEnum from roborock.util import get_next_int, get_timestamp @@ -37,8 +36,8 @@ class RoborockDataProtocol(RoborockEnum): OFFLINE_STATUS = 135 @classmethod - def _missing_(cls: type[RoborockEnum], key) -> RoborockEnum: - raise ValueError("%s not a valid key for Data Protocol", key) + def _missing_(cls: type[Self], key) -> Self: + raise ValueError(f"{key} not a valid key for Data Protocol") class RoborockDyadDataProtocol(RoborockEnum): diff --git a/roborock/roborock_typing.py b/roborock/roborock_typing.py index 418f350c..71251324 100644 --- a/roborock/roborock_typing.py +++ b/roborock/roborock_typing.py @@ -1,7 +1,6 @@ -from __future__ import annotations - from dataclasses import dataclass, field from enum import Enum, StrEnum +from typing import Self from .data import ( CleanRecord, @@ -368,7 +367,7 @@ def __post_init__(self) -> None: ): self.dust_collection_mode_name = self.dock_summary.dust_collection_mode.mode.name - def update(self, device_prop: DeviceProp) -> None: + def update(self, device_prop: Self) -> None: if device_prop.status: self.status = device_prop.status if device_prop.clean_summary: diff --git a/roborock/util.py b/roborock/util.py index 8679511e..481759ec 100644 --- a/roborock/util.py +++ b/roborock/util.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import logging import math import time diff --git a/roborock/web_api.py b/roborock/web_api.py index 141cbeb5..a76d14c5 100644 --- a/roborock/web_api.py +++ b/roborock/web_api.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import base64 import hashlib import hmac diff --git a/tests/data/b01_q10/test_b01_q10_code_mappings.py b/tests/data/b01_q10/test_b01_q10_code_mappings.py new file mode 100644 index 00000000..63fdd14d --- /dev/null +++ b/tests/data/b01_q10/test_b01_q10_code_mappings.py @@ -0,0 +1,35 @@ +"""Test cases for B01 Q10 code mappings.""" + +import pytest + +from roborock.data.b01_q10 import YXDeviceState + + +@pytest.mark.parametrize( + "state, string", + [ + (YXDeviceState.UNKNOWN, "unknown"), + (YXDeviceState.IDLE, "idle"), + (YXDeviceState.CHARGING, "charging"), + (YXDeviceState.CLEANING, "cleaning"), + (YXDeviceState.SLEEPING, "sleeping"), + (YXDeviceState.UPDATING, "updating"), + (YXDeviceState.RETURNING_HOME, "returning_home"), + ], +) +def test_q10_status_values_are_canonical(state: YXDeviceState, string: str) -> None: + """Q10 status enum values should expose canonical names.""" + assert state.value == string + + +@pytest.mark.parametrize( + "code, expected_state", + [ + (5, YXDeviceState.CLEANING), + (8, YXDeviceState.CHARGING), + (14, YXDeviceState.UPDATING), + ], +) +def test_q10_status_codes_map_to_canonical_values(code: int, expected_state: YXDeviceState) -> None: + """Code-based mapping should return canonical status values.""" + assert YXDeviceState.from_code(code) is expected_state diff --git a/tests/data/test_code_mappings.py b/tests/data/test_code_mappings.py index 5cd48fa2..875d6c8c 100644 --- a/tests/data/test_code_mappings.py +++ b/tests/data/test_code_mappings.py @@ -5,7 +5,7 @@ import pytest from roborock import HomeDataProduct, RoborockCategory -from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, YXCleanType def test_from_code() -> None: @@ -52,6 +52,30 @@ def test_invalid_from_value() -> None: B01_Q10_DP.from_value("invalid_value") +@pytest.mark.parametrize( + "input, expected", + [ + ("START_CLEAN", B01_Q10_DP.START_CLEAN), + ("start_clean", B01_Q10_DP.START_CLEAN), + ("dpStartClean", B01_Q10_DP.START_CLEAN), + (201, B01_Q10_DP.START_CLEAN), + ("PAUSE", B01_Q10_DP.PAUSE), + ("pause", B01_Q10_DP.PAUSE), + ("dpPause", B01_Q10_DP.PAUSE), + (204, B01_Q10_DP.PAUSE), + ("STOP", B01_Q10_DP.STOP), + ("stop", B01_Q10_DP.STOP), + ("dpStop", B01_Q10_DP.STOP), + (206, B01_Q10_DP.STOP), + ("invalid_value", None), + (999999, None), + ], +) +def test_from_any_optional(input: str | int, expected: B01_Q10_DP | None) -> None: + """Test from_any_optional method.""" + assert B01_Q10_DP.from_any_optional(input) == expected + + def test_homedata_product_unknown_category(): """Test that HomeDataProduct can handle unknown categories.""" data = { @@ -65,3 +89,17 @@ def test_homedata_product_unknown_category(): product = HomeDataProduct.from_dict(data) assert product.id == "unknown_cat_id" assert product.category == RoborockCategory.UNKNOWN + + +@pytest.mark.parametrize( + ("readable_value", "expected_clean_type"), + [ + ("vac_and_mop", YXCleanType.VAC_AND_MOP), + ("vacuum", YXCleanType.VACUUM), + ("mop", YXCleanType.MOP), + ], +) +def test_yx_clean_type_from_value_readable_values(readable_value: str, expected_clean_type: YXCleanType) -> None: + """Test YXCleanType accepts canonical readable values.""" + assert YXCleanType.from_value(readable_value) is expected_clean_type + assert expected_clean_type.value == readable_value diff --git a/tests/devices/__snapshots__/test_v1_device.ambr b/tests/devices/__snapshots__/test_v1_device.ambr index 2ad1cb7e..8da5d44c 100644 --- a/tests/devices/__snapshots__/test_v1_device.ambr +++ b/tests/devices/__snapshots__/test_v1_device.ambr @@ -1,6 +1,6 @@ # serializer version: 1 # name: test_device_trait_command_parsing[clean_summary] - CleanSummaryTrait(clean_area=24258125000, clean_count=296, clean_time=1442559, command=, dust_collection_count=None, last_clean_record=CleanRecord(area=81122500, avoid_count=None, begin=1738864366, begin_datetime=datetime.datetime(2025, 2, 6, 17, 52, 46, tzinfo=datetime.timezone.utc), clean_type=None, complete=None, duration=4358, dust_collection_status=None, end=1738868964, end_datetime=datetime.datetime(2025, 2, 6, 19, 9, 24, tzinfo=datetime.timezone.utc), error=None, finish_reason=None, map_flag=None, square_meter_area=81.1, start_type=None, wash_count=None), last_clean_t=None, records=[1756848207, 1754930385, 1753203976, 1752183435, 1747427370, 1746204046, 1745601543, 1744387080, 1743528522, 1742489154, 1741022299, 1740433682, 1739902516, 1738875106, 1738864366, 1738620067, 1736873889, 1736197544, 1736121269, 1734458038], square_meter_clean_area=24258.1) + CleanSummaryTrait(clean_area=24258125000, clean_count=296, clean_record_converter=CleanRecordConverter, clean_time=1442559, command=, converter=CleanSummaryConverter, dust_collection_count=None, last_clean_record=CleanRecord(area=81122500, avoid_count=None, begin=1738864366, begin_datetime=datetime.datetime(2025, 2, 6, 17, 52, 46, tzinfo=datetime.timezone.utc), clean_type=None, complete=None, duration=4358, dust_collection_status=None, end=1738868964, end_datetime=datetime.datetime(2025, 2, 6, 19, 9, 24, tzinfo=datetime.timezone.utc), error=None, finish_reason=None, map_flag=None, square_meter_area=81.1, start_type=None, wash_count=None), last_clean_t=None, records=[1756848207, 1754930385, 1753203976, 1752183435, 1747427370, 1746204046, 1745601543, 1744387080, 1743528522, 1742489154, 1741022299, 1740433682, 1739902516, 1738875106, 1738864366, 1738620067, 1736873889, 1736197544, 1736121269, 1734458038], square_meter_clean_area=24258.1) # --- # name: test_device_trait_command_parsing[clean_summary].1 dict({ @@ -870,7 +870,7 @@ }) # --- # name: test_device_trait_command_parsing[status] - StatusTrait(adbumper_status=None, auto_dust_collection=None, avoid_count=None, back_type=None, battery=100, camera_status=None, charge_status=None, clean_area=91287500, clean_fluid_status=None, clean_percent=None, clean_time=5405, clear_water_box_status=None, collision_avoid_status=None, command=, common_status=None, corner_clean_mode=None, current_map=0, debug_mode=None, dirty_water_box_status=None, distance_off=0, dnd_enabled=1, dock_cool_fan_status=None, dock_error_status=None, dock_type=None, dry_status=None, dss=None, dust_bag_status=None, dust_collection_status=None, error_code=, error_code_name='none', fan_power=106, fan_speed_mapping={101: 'quiet', 102: 'balanced', 103: 'turbo', 104: 'max', 108: 'max_plus', 105: 'off', 106: 'custom'}, fan_speed_name='custom', fan_speed_options=[, , , , , , ], hatch_door_status=None, home_sec_enable_password=None, home_sec_status=None, in_cleaning=, in_fresh_state=1, in_returning=0, in_warmup=None, is_exploring=None, is_locating=0, kct=None, lab_status=1, last_clean_t=None, lock_status=0, map_present=1, map_status=3, mop_forbidden_enable=0, mop_mode=None, mop_route_mapping={300: 'standard', 301: 'deep', 302: 'custom'}, mop_route_name=None, mop_route_options=[, , ], msg_seq=515, msg_ver=2, rdt=None, repeat=None, replenish_mode=None, rss=None, square_meter_clean_area=91.3, state=, state_name='charging', subdivision_sets=None, switch_map_mode=None, unsave_map_flag=0, unsave_map_reason=4, wash_phase=None, wash_ready=None, wash_status=None, water_box_carriage_status=0, water_box_filter_status=None, water_box_mode=204, water_box_status=0, water_mode_mapping={200: 'off', 201: 'mild', 202: 'standard', 203: 'intense', 207: 'custom_water_flow', 204: 'custom'}, water_mode_name='custom', water_mode_options=[, , , , , ], water_shortage_status=None) + StatusTrait(adbumper_status=None, auto_dust_collection=None, avoid_count=None, back_type=None, battery=100, camera_status=None, charge_status=None, clean_area=91287500, clean_fluid_status=None, clean_percent=None, clean_time=5405, clear_water_box_status=None, collision_avoid_status=None, command=, common_status=None, converter=DefaultConverter, corner_clean_mode=None, current_map=0, debug_mode=None, dirty_water_box_status=None, distance_off=0, dnd_enabled=1, dock_cool_fan_status=None, dock_error_status=None, dock_type=None, dry_status=None, dss=None, dust_bag_status=None, dust_collection_status=None, error_code=, error_code_name='none', fan_power=106, fan_speed_mapping={101: 'quiet', 102: 'balanced', 103: 'turbo', 104: 'max', 108: 'max_plus', 105: 'off', 106: 'custom'}, fan_speed_name='custom', fan_speed_options=[, , , , , , ], hatch_door_status=None, home_sec_enable_password=None, home_sec_status=None, in_cleaning=, in_fresh_state=1, in_returning=0, in_warmup=None, is_exploring=None, is_locating=0, kct=None, lab_status=1, last_clean_t=None, lock_status=0, map_present=1, map_status=3, mop_forbidden_enable=0, mop_mode=None, mop_route_mapping={300: 'standard', 301: 'deep', 302: 'custom'}, mop_route_name=None, mop_route_options=[, , ], msg_seq=515, msg_ver=2, rdt=None, repeat=None, replenish_mode=None, rss=None, square_meter_clean_area=91.3, state=, state_name='charging', subdivision_sets=None, switch_map_mode=None, unsave_map_flag=0, unsave_map_reason=4, wash_phase=None, wash_ready=None, wash_status=None, water_box_carriage_status=0, water_box_filter_status=None, water_box_mode=204, water_box_status=0, water_mode_mapping={200: 'off', 201: 'mild', 202: 'standard', 203: 'intense', 207: 'custom_water_flow', 204: 'custom'}, water_mode_name='custom', water_mode_options=[, , , , , ], water_shortage_status=None) # --- # name: test_device_trait_command_parsing[status].1 dict({ diff --git a/tests/devices/traits/b01/q10/test_status.py b/tests/devices/traits/b01/q10/test_status.py index e295e890..9e526c7e 100644 --- a/tests/devices/traits/b01/q10/test_status.py +++ b/tests/devices/traits/b01/q10/test_status.py @@ -96,10 +96,10 @@ async def test_status_trait_streaming( message_queue.put_nowait(message) # Wait for the update - await wait_for_attribute_value(q10_api.status, "status", YXDeviceState.CHARGING_STATE) + await wait_for_attribute_value(q10_api.status, "status", YXDeviceState.CHARGING) # Verify trait attributes are updated - assert q10_api.status.status == YXDeviceState.CHARGING_STATE + assert q10_api.status.status == YXDeviceState.CHARGING assert q10_api.status.clean_task_type == YXDeviceCleanTask.IDLE @@ -112,11 +112,15 @@ async def test_status_trait_refresh( assert q10_api.status.battery is None assert q10_api.status.status is None assert q10_api.status.fan_level is None + assert q10_api.status.total_clean_count is None + assert q10_api.status.main_brush_life is None + assert q10_api.status.cleaning_progress is None + assert q10_api.status.fault is None # Mock the response to refresh # battery (122) = 100 # status (121) = 8 (CHARGING_STATE) - # fun_level (123) = 2 (NORMAL) + # fan_level (123) = 2 (BALANCED) message = build_message(TESTDATA_DP_REQUEST_DPS) # Send a refresh command @@ -138,8 +142,17 @@ async def test_status_trait_refresh( # Verify trait attributes are updated assert q10_api.status.battery == 100 - assert q10_api.status.status == YXDeviceState.CHARGING_STATE - assert q10_api.status.fan_level == YXFanLevel.NORMAL + assert q10_api.status.status == YXDeviceState.CHARGING + assert q10_api.status.fan_level == YXFanLevel.BALANCED + assert q10_api.status.total_clean_area == 0 + assert q10_api.status.total_clean_count == 0 + assert q10_api.status.total_clean_time == 0 + assert q10_api.status.main_brush_life == 0 + assert q10_api.status.side_brush_life == 0 + assert q10_api.status.filter_life == 0 + assert q10_api.status.sensor_life == 0 + assert q10_api.status.cleaning_progress == 100 + assert q10_api.status.fault == 0 def test_status_trait_update_listener(q10_api: Q10PropertiesApi) -> None: diff --git a/tests/devices/traits/b01/q10/test_vacuum.py b/tests/devices/traits/b01/q10/test_vacuum.py index c8bdb3a4..499b9bbc 100644 --- a/tests/devices/traits/b01/q10/test_vacuum.py +++ b/tests/devices/traits/b01/q10/test_vacuum.py @@ -34,8 +34,8 @@ def vacuumm_fixture(q10_api: Q10PropertiesApi) -> VacuumTrait: (lambda x: x.stop_clean(), {"206": {}}), (lambda x: x.return_to_dock(), {"203": {}}), (lambda x: x.empty_dustbin(), {"203": 2}), - (lambda x: x.set_clean_mode(YXCleanType.BOTH_WORK), {"137": 1}), - (lambda x: x.set_fan_level(YXFanLevel.NORMAL), {"123": 2}), + (lambda x: x.set_clean_mode(YXCleanType.VAC_AND_MOP), {"137": 1}), + (lambda x: x.set_fan_level(YXFanLevel.BALANCED), {"123": 2}), ], ) async def test_vacuum_commands( diff --git a/tests/devices/traits/b01/q7/conftest.py b/tests/devices/traits/b01/q7/conftest.py index 5dc476f6..4f55ba0f 100644 --- a/tests/devices/traits/b01/q7/conftest.py +++ b/tests/devices/traits/b01/q7/conftest.py @@ -5,6 +5,7 @@ import pytest +from roborock.data import HomeDataDevice, HomeDataProduct, RoborockCategory from roborock.devices.traits.b01.q7 import Q7PropertiesApi from tests.fixtures.channel_fixtures import FakeChannel @@ -16,9 +17,30 @@ def fake_channel_fixture() -> FakeChannel: return FakeChannel() +@pytest.fixture(name="product") +def product_fixture() -> HomeDataProduct: + return HomeDataProduct( + id="product-id-q7", + name="Roborock Q7", + model="roborock.vacuum.sc05", + category=RoborockCategory.VACUUM, + ) + + +@pytest.fixture(name="device") +def device_fixture() -> HomeDataDevice: + return HomeDataDevice( + duid="abc123", + name="Q7", + local_key="key123key123key1", + product_id="product-id-q7", + sn="testsn012345", + ) + + @pytest.fixture(name="q7_api") -def q7_api_fixture(fake_channel: FakeChannel) -> Q7PropertiesApi: - return Q7PropertiesApi(fake_channel) # type: ignore[arg-type] +def q7_api_fixture(fake_channel: FakeChannel, device: HomeDataDevice, product: HomeDataProduct) -> Q7PropertiesApi: + return Q7PropertiesApi(fake_channel, device=device, product=product) # type: ignore[arg-type] @pytest.fixture(name="expected_msg_id", autouse=True) @@ -28,6 +50,7 @@ def next_message_id_fixture() -> Generator[int, None, None]: We pick an arbitrary number, but just need it to ensure we can craft a fake response with the message id matched to the outgoing RPC. """ + expected_msg_id = math.floor(time.time()) # Patch get_next_int to return our expected msg_id so the channel waits for it diff --git a/tests/devices/traits/b01/q7/test_map_content.py b/tests/devices/traits/b01/q7/test_map_content.py new file mode 100644 index 00000000..c0b62a4f --- /dev/null +++ b/tests/devices/traits/b01/q7/test_map_content.py @@ -0,0 +1,79 @@ +from typing import cast +from unittest.mock import patch + +import pytest +from vacuum_map_parser_base.map_data import MapData + +from roborock.devices.traits.b01.q7 import Q7PropertiesApi +from roborock.devices.transport.mqtt_channel import MqttChannel +from roborock.exceptions import RoborockException +from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol +from tests.fixtures.channel_fixtures import FakeChannel + +from . import B01MessageBuilder + + +async def test_q7_map_content_refresh_populates_cached_values( + q7_api: Q7PropertiesApi, + fake_channel: FakeChannel, + message_builder: B01MessageBuilder, +): + fake_channel.response_queue.append(message_builder.build({"map_list": [{"id": 1772093512, "cur": True}]})) + fake_channel.response_queue.append( + RoborockMessage( + protocol=RoborockMessageProtocol.MAP_RESPONSE, + payload=b"raw-map-payload", + version=b"B01", + seq=message_builder.seq + 1, + ) + ) + + dummy_map_data = MapData() + with patch( + "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", + return_value=type("X", (), {"image_content": b"pngbytes", "map_data": dummy_map_data})(), + ) as parse: + await q7_api.map_content.refresh() + + assert q7_api.map_content.image_content == b"pngbytes" + assert q7_api.map_content.map_data is dummy_map_data + assert q7_api.map_content.raw_api_response == b"raw-map-payload" + + parse.assert_called_once() + + +def test_q7_map_content_parse_errors_cleanly(q7_api: Q7PropertiesApi): + with patch("roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", side_effect=ValueError("boom")): + with pytest.raises(RoborockException, match="Failed to parse B01 map data"): + q7_api.map_content.parse_map_content(b"raw") + + +def test_q7_map_content_preserves_specific_roborock_errors(q7_api: Q7PropertiesApi): + with patch( + "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", + side_effect=RoborockException("Specific decoder failure"), + ): + with pytest.raises(RoborockException, match="Specific decoder failure"): + q7_api.map_content.parse_map_content(b"raw") + + +def test_q7_map_content_requires_metadata_at_init(fake_channel: FakeChannel): + from roborock.data import HomeDataDevice, HomeDataProduct, RoborockCategory + + with pytest.raises(ValueError, match="requires device serial number and product model metadata"): + Q7PropertiesApi( + cast(MqttChannel, fake_channel), + device=HomeDataDevice( + duid="abc123", + name="Q7", + local_key="key123key123key1", + product_id="product-id-q7", + sn=None, + ), + product=HomeDataProduct( + id="product-id-q7", + name="Roborock Q7", + model="roborock.vacuum.sc05", + category=RoborockCategory.VACUUM, + ), + ) diff --git a/tests/devices/traits/v1/test_status.py b/tests/devices/traits/v1/test_status.py index 21cd5844..a308dbca 100644 --- a/tests/devices/traits/v1/test_status.py +++ b/tests/devices/traits/v1/test_status.py @@ -63,7 +63,7 @@ async def test_refresh_status_invalid_format(status_trait: StatusTrait, mock_rpc """Test that invalid response format raises ValueError.""" mock_rpc_channel.send_command.return_value = "invalid" - with pytest.raises(ValueError, match="Unexpected status format"): + with pytest.raises(ValueError, match="Unexpected StatusV2 response format"): await status_trait.refresh() diff --git a/tests/map/test_b01_map_parser.py b/tests/map/test_b01_map_parser.py new file mode 100644 index 00000000..6c91dba0 --- /dev/null +++ b/tests/map/test_b01_map_parser.py @@ -0,0 +1,175 @@ +import base64 +import gzip +import hashlib +import io +import struct +import zlib +from pathlib import Path + +import pytest +from Crypto.Cipher import AES +from Crypto.Util.Padding import pad +from PIL import Image + +from roborock.exceptions import RoborockException +from roborock.map.b01_map_parser import B01MapParser, _parse_scmap_payload + +FIXTURE = Path(__file__).resolve().parent / "testdata" / "raw-mqtt-map301.bin.inflated.bin.gz" + + +def _derive_map_key(serial: str, model: str) -> bytes: + model_suffix = model.split(".")[-1] + model_key = (model_suffix + "0" * 16)[:16].encode() + material = f"{serial}+{model_suffix}+{serial}".encode() + encrypted = AES.new(model_key, AES.MODE_ECB).encrypt(pad(material, AES.block_size)) + md5 = hashlib.md5(base64.b64encode(encrypted), usedforsecurity=False).hexdigest() + return md5[8:24].encode() + + +def _encode_varint(value: int) -> bytes: + encoded = bytearray() + while True: + to_write = value & 0x7F + value >>= 7 + if value: + encoded.append(to_write | 0x80) + else: + encoded.append(to_write) + return bytes(encoded) + + +def _field_varint(field_no: int, value: int) -> bytes: + return _encode_varint((field_no << 3) | 0) + _encode_varint(value) + + +def _field_len(field_no: int, value: bytes) -> bytes: + return _encode_varint((field_no << 3) | 2) + _encode_varint(len(value)) + value + + +def _field_float32(field_no: int, value: float) -> bytes: + return _encode_varint((field_no << 3) | 5) + struct.pack(" None: + serial = "testsn012345" + model = "roborock.vacuum.sc05" + inflated = gzip.decompress(FIXTURE.read_bytes()) + + compressed = zlib.compress(inflated) + map_key = _derive_map_key(serial, model) + encrypted = AES.new(map_key, AES.MODE_ECB).encrypt(pad(compressed.hex().encode(), AES.block_size)) + payload = base64.b64encode(encrypted) + + parser = B01MapParser() + parsed = parser.parse(payload, serial=serial, model=model) + + assert parsed.image_content is not None + assert parsed.image_content.startswith(b"\x89PNG\r\n\x1a\n") + assert parsed.map_data is not None + + # The fixture includes 10 rooms with names room1..room10. + assert parsed.map_data.additional_parameters["room_names"] == { + 10: "room1", + 11: "room2", + 12: "room3", + 13: "room4", + 14: "room5", + 15: "room6", + 16: "room7", + 17: "room8", + 18: "room9", + 19: "room10", + } + + # Image should be scaled by default. + img = Image.open(io.BytesIO(parsed.image_content)) + assert img.size == (340 * 4, 300 * 4) + + +def test_b01_scmap_parser_maps_observed_schema_fields() -> None: + room_name_post = _field_float32(1, 11.25) + _field_float32(2, 22.5) + room_one = b"".join( + [ + _field_varint(1, 42), + _field_len(2, b"Kitchen"), + _field_varint(5, 1), + _field_len(8, room_name_post), + _field_varint(10, 7), + _field_varint(12, 9), + ] + ) + room_two = b"".join([_field_varint(1, 99), _field_varint(5, 0)]) + + boundary_info = b"".join( + [ + _field_len(1, b"md5"), + _field_varint(2, 10), + _field_varint(3, 20), + _field_varint(4, 30), + _field_varint(5, 40), + ] + ) + map_ext_info = b"".join( + [ + _field_varint(1, 100), + _field_varint(2, 200), + _field_varint(3, 1), + _field_varint(8, 3), + _field_len(7, boundary_info), + ] + ) + map_head = b"".join( + [ + _field_varint(1, 7), + _field_varint(2, 2), + _field_varint(3, 2), + _field_float32(4, 1.5), + _field_float32(5, 2.5), + _field_float32(6, 3.5), + _field_float32(7, 4.5), + _field_float32(8, 0.05), + ] + ) + map_data = _field_len(1, zlib.compress(bytes([0, 127, 128, 128]))) + payload = b"".join( + [ + _field_varint(1, 1), + _field_len(2, map_ext_info), + _field_len(3, map_head), + _field_len(4, map_data), + _field_len(12, room_one), + _field_len(12, room_two), + ] + ) + + parsed = _parse_scmap_payload(payload) + + assert parsed.mapType == 1 + assert parsed.HasField("mapExtInfo") + assert parsed.mapExtInfo.taskBeginDate == 100 + assert parsed.mapExtInfo.mapUploadDate == 200 + assert parsed.mapExtInfo.HasField("boudaryInfo") + assert parsed.mapExtInfo.boudaryInfo.vMaxY == 40 + assert parsed.HasField("mapHead") + assert parsed.mapHead.mapHeadId == 7 + assert parsed.mapHead.sizeX == 2 + assert parsed.mapHead.sizeY == 2 + assert parsed.mapHead.resolution == pytest.approx(0.05) + assert parsed.HasField("mapData") + assert parsed.mapData.HasField("mapData") + assert zlib.decompress(parsed.mapData.mapData) == bytes([0, 127, 128, 128]) + assert parsed.roomDataInfo[0].roomId == 42 + assert parsed.roomDataInfo[0].roomName == "Kitchen" + assert parsed.roomDataInfo[0].HasField("roomNamePost") + assert parsed.roomDataInfo[0].roomNamePost.x == pytest.approx(11.25) + assert parsed.roomDataInfo[0].roomNamePost.y == pytest.approx(22.5) + assert parsed.roomDataInfo[0].colorId == 7 + assert parsed.roomDataInfo[0].global_seq == 9 + assert parsed.roomDataInfo[1].roomId == 99 + assert not parsed.roomDataInfo[1].HasField("roomName") + + +def test_b01_map_parser_rejects_invalid_payload() -> None: + parser = B01MapParser() + with pytest.raises(RoborockException, match="Failed to decode B01 map payload"): + parser.parse(b"not a map", serial="testsn012345", model="roborock.vacuum.sc05") diff --git a/tests/map/testdata/raw-mqtt-map301.bin.inflated.bin.gz b/tests/map/testdata/raw-mqtt-map301.bin.inflated.bin.gz new file mode 100644 index 00000000..63b67ddb Binary files /dev/null and b/tests/map/testdata/raw-mqtt-map301.bin.inflated.bin.gz differ diff --git a/tests/protocols/test_b01_q10_protocol.py b/tests/protocols/test_b01_q10_protocol.py index 110b20f0..62ee2c27 100644 --- a/tests/protocols/test_b01_q10_protocol.py +++ b/tests/protocols/test_b01_q10_protocol.py @@ -97,7 +97,7 @@ def test_decode_unknown_dps_code() -> None: (B01_Q10_DP.REQUEST_DPS, {}), (B01_Q10_DP.REQUEST_DPS, None), (B01_Q10_DP.START_CLEAN, {"cmd": 1}), - (B01_Q10_DP.WATER_LEVEL, YXWaterLevel.MIDDLE.code), + (B01_Q10_DP.WATER_LEVEL, YXWaterLevel.MEDIUM.code), ], ) def test_encode_mqtt_payload(command: B01_Q10_DP, params: dict[str, Any], snapshot) -> None: diff --git a/uv.lock b/uv.lock index 57b98421..4cf57d0d 100644 --- a/uv.lock +++ b/uv.lock @@ -1320,7 +1320,7 @@ wheels = [ [[package]] name = "python-roborock" -version = "4.22.0" +version = "5.0.0" source = { editable = "." } dependencies = [ { name = "aiohttp" },