From 7f8b8380da71e30822f710e252fd1c21061fab7e Mon Sep 17 00:00:00 2001
From: Paul Frybarger <pfrybar@frymail.net>
Date: Fri, 21 May 2021 13:38:31 +0100
Subject: [PATCH] Add zone support for Dyson 360 Heurist

---
 libdyson/cloud/cloud_360_heurist.py | 62 +++++++++++++++++++++++++++++
 libdyson/dyson_360_heurist.py       | 26 +++++++++++-
 libdyson/utils.py                   |  9 ++++-
 requirements.txt                    |  1 +
 4 files changed, 95 insertions(+), 3 deletions(-)
 create mode 100644 libdyson/cloud/cloud_360_heurist.py

diff --git a/libdyson/cloud/cloud_360_heurist.py b/libdyson/cloud/cloud_360_heurist.py
new file mode 100644
index 0000000..88b016d
--- /dev/null
+++ b/libdyson/cloud/cloud_360_heurist.py
@@ -0,0 +1,62 @@
+"""Dyson 360 Heurist cloud client."""
+
+from datetime import datetime
+from typing import List
+
+import attr
+from dateutil import parser
+
+from .cloud_device import DysonCloudDevice
+
+
+@attr.s(auto_attribs=True, frozen=True)
+class Zone:
+    """Represent a zone within a persistent map."""
+
+    id: str
+    name: str
+    icon: str
+    area: float  # In square meters
+
+    @classmethod
+    def from_raw(cls, raw: dict):
+        return cls(
+            raw["id"],
+            raw["name"],
+            raw["icon"],
+            raw["area"],
+        )
+
+
+@attr.s(auto_attribs=True, frozen=True)
+class PersistentMap:
+    """Represent a persistent map created by the user."""
+
+    id: str
+    name: str
+    last_visited: datetime  # UTC
+    zones_definition_last_updated_date: datetime  # UTC
+    zones: List[Zone]
+
+    @classmethod
+    def from_raw(cls, raw: dict):
+        """Parse raw data from cloud API."""
+        return cls(
+            raw["id"],
+            raw["name"],
+            parser.isoparse(raw["lastVisited"]),
+            parser.isoparse(raw["zonesDefinitionLastUpdatedDate"]),
+            [Zone.from_raw(raw) for raw in raw["zones"]],
+        )
+
+
+class DysonCloud360Heurist(DysonCloudDevice):
+    """Dyson 360 Heurist cloud client."""
+
+    def get_persistent_maps(self) -> List[PersistentMap]:
+        """Get the persistent maps from the cloud."""
+        response = self._account.request(
+            "GET",
+            f"/v1/app/{self._serial}/persistent-map-metadata",
+        )
+        return [PersistentMap.from_raw(raw) for raw in response.json()]
diff --git a/libdyson/dyson_360_heurist.py b/libdyson/dyson_360_heurist.py
index 70c8911..4424916 100644
--- a/libdyson/dyson_360_heurist.py
+++ b/libdyson/dyson_360_heurist.py
@@ -1,9 +1,11 @@
 """Dyson 360 Heurist vacuum robot."""
 
-from typing import Optional
+from typing import Optional, List
 
 from .const import DEVICE_TYPE_360_HEURIST, CleaningMode, VacuumHeuristPowerMode
 from .dyson_vacuum_device import DysonVacuumDevice
+from .utils import mqtt_time_from_datetime
+from .cloud.cloud_360_heurist import PersistentMap
 
 
 class Dyson360Heurist(DysonVacuumDevice):
@@ -56,6 +58,28 @@ def start_all_zones(self) -> None:
             "START", {"cleaningMode": "global", "fullCleanType": "immediate"}
         )
 
+    def start_zones(self, persistent_map: PersistentMap, zone_names: List[str]) -> None:
+        """Start cleaning of specific zone(s)."""
+        zone_ids = []
+        for zone_name in zone_names:
+            zone = next(filter(lambda z: z.name == zone_name, persistent_map.zones), None)
+            if zone is None:
+                raise ValueError("Invalid zone %s", zone_name)
+            zone_ids.append(zone.id)
+
+        data = {
+            "cleaningMode": "zoneConfigured",
+            "fullCleanType": "immediate",
+            "cleaningProgramme": {
+                "persistentMapId": persistent_map.id,
+                "zonesDefinitionLastUpdatedDate": mqtt_time_from_datetime(persistent_map.zones_definition_last_updated_date),
+                "orderedZones": [],
+                "unorderedZones": zone_ids,
+            },
+        }
+
+        self._send_command("START", data)
+
     def set_default_power_mode(self, power_mode: VacuumHeuristPowerMode) -> None:
         """Set default power mode."""
         self._send_command(
diff --git a/libdyson/utils.py b/libdyson/utils.py
index 85f1a16..fe4f03f 100644
--- a/libdyson/utils.py
+++ b/libdyson/utils.py
@@ -3,7 +3,7 @@
 import base64
 import hashlib
 import re
-import time
+from datetime import datetime
 from typing import Tuple
 
 from .const import DEVICE_TYPE_360_EYE
@@ -12,7 +12,12 @@
 
 def mqtt_time():
     """Return current time string for mqtt messages."""
-    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+    return mqtt_time_from_datetime(datetime.utcnow())
+
+
+def mqtt_time_from_datetime(dt: datetime):
+    """Return time string for mqtt messages."""
+    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")
 
 
 def get_credential_from_wifi_password(wifi_password: str) -> str:
diff --git a/requirements.txt b/requirements.txt
index caeee57..f17357b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,3 +3,4 @@ cryptography>=3.1
 requests
 zeroconf
 attrs
+python-dateutil