roborock.devices.traits.traits_mixin

 1"""Holds device traits mixin and related code.
 2
 3This holds the TraitsMixin class, which is used to provide accessors for
 4various device traits. Each trait is a class that encapsulates a specific
 5set of functionality for a device, such as controlling a vacuum or a mop.
 6
 7The TraitsMixin holds traits across all protocol types. A trait is supported
 8if it is non-None.
 9"""
10
11from dataclasses import dataclass, fields
12from typing import get_args, get_origin
13
14from . import Trait, a01, b01, v1
15
16__all__ = [
17    "TraitsMixin",
18]
19
20
21@dataclass(init=False)
22class TraitsMixin:
23    """Mixin to provide trait accessors."""
24
25    v1_properties: v1.PropertiesApi | None = None
26    """V1 properties trait, if supported."""
27
28    dyad: a01.DyadApi | None = None
29    """Dyad API, if supported."""
30
31    zeo: a01.ZeoApi | None = None
32    """Zeo API, if supported."""
33
34    b01_q7_properties: b01.Q7PropertiesApi | None = None
35    """B01 Q7 properties trait, if supported."""
36
37    def __init__(self, trait: Trait) -> None:
38        """Initialize the TraitsMixin with the given trait.
39
40        This will populate the appropriate trait attributes based on the types
41        of the traits provided.
42        """
43        for item in fields(self):
44            trait_type = _get_trait_type(item)
45            if trait_type is type(trait):
46                setattr(self, item.name, trait)
47                break
48
49
50def _get_trait_type(item) -> type[Trait]:
51    """Get the trait type from a dataclass field."""
52    if get_origin(item.type) is None:
53        raise ValueError(f"Trait {item.name} is not an optional type")
54    if (args := get_args(item.type)) is None:
55        raise ValueError(f"Trait {item.name} is not an optional type")
56    if len(args) != 2 or args[1] is not type(None):
57        raise ValueError(f"Trait {item.name} is not an optional type")
58    trait_type = args[0]
59    if not issubclass(trait_type, Trait):
60        raise ValueError(f"Trait {item.name} is not a Trait subclass")
61    return trait_type