example.py

"""Example script demonstrating how to connect to Roborock devices and print their status."""

import asyncio
import dataclasses
import json
import pathlib
from typing import Any

from roborock.devices.device_manager import UserParams, create_device_manager
from roborock.devices.file_cache import FileCache, load_value, store_value
from roborock.web_api import RoborockApiClient

# We typically store the login credentials/information separately from other cached data.
USER_PARAMS_PATH = pathlib.Path.home() / ".cache" / "roborock-user-params.pkl"

# Device connection information is cached to speed up future connections.
CACHE_PATH = pathlib.Path.home() / ".cache" / "roborock-cache-data.pkl"


async def login_flow() -> UserParams:
    """Perform the login flow to obtain UserData from the web API."""
    username = input("Email: ")
    web_api = RoborockApiClient(username=username)
    print("Requesting login code sent to email...")
    await web_api.request_code()
    code = input("Code: ")
    user_data = await web_api.code_login(code)
    # We store the base_url to avoid future discovery calls.
    base_url = await web_api.base_url
    return UserParams(
        username=username,
        user_data=user_data,
        base_url=base_url,
    )


async def get_or_create_session() -> UserParams:
    """Initialize the session by logging in if necessary."""
    user_params = await load_value(USER_PARAMS_PATH)
    if user_params is None:
        print("No cached login data found, please login.")
        user_params = await login_flow()
        print("Login successful, caching login data...")
        await store_value(USER_PARAMS_PATH, user_params)
        print(f"Cached login data to {USER_PARAMS_PATH}.")
    return user_params


def remove_none_values(data: dict[str, Any]) -> dict[str, Any]:
    return {k: v for k, v in data.items() if v is not None}


async def main():
    user_params = await get_or_create_session()
    cache = FileCache(CACHE_PATH)

    # Create a device manager that can discover devices.
    device_manager = await create_device_manager(user_params, cache=cache)
    devices = await device_manager.get_devices()

    # Get all vacuum devices that support the v1 PropertiesApi
    device_results = []
    for device in devices:
        if not device.v1_properties:
            continue

        # Refresh the current device status
        status_trait = device.v1_properties.status
        await status_trait.refresh()

        # Print the device status as JSON
        device_results.append(
            {
                "device": device.name,
                "status": remove_none_values(dataclasses.asdict(status_trait)),
            }
        )

    print(json.dumps(device_results, indent=2))

    await cache.flush()


if __name__ == "__main__":
    asyncio.run(main())