python-roborock 2.45.0__tar.gz → 2.47.0__tar.gz

{python_roborock-2.45.0 → python_roborock-2.47.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: python-roborock
-Version: 2.45.0
+Version: 2.47.0
 Summary: A package to control Roborock vacuums.
 Home-page: https://github.com/humbertogontijo/python-roborock
 License: GPL-3.0-only

{python_roborock-2.45.0 → python_roborock-2.47.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.45.0"
+version = "2.47.0"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"

{python_roborock-2.45.0 → python_roborock-2.47.0}/roborock/api.py

@@ -18,12 +18,11 @@ from .exceptions import (
 from .roborock_future import RoborockFuture
 from .roborock_message import (
     RoborockMessage,
-    RoborockMessageProtocol,
 )
 from .util import get_next_int
 
 _LOGGER = logging.getLogger(__name__)
-KEEPALIVE = 60
+KEEPALIVE = 70
 
 
 class RoborockClient(ABC):
@@ -78,12 +77,6 @@ class RoborockClient(ABC):
             return False
         return True
 
-    async def validate_connection(self) -> None:
-        if not self.should_keepalive():
-            self._logger.info("Resetting Roborock connection due to keepalive timeout")
-            await self.async_disconnect()
-        await self.async_connect()
-
     async def _wait_response(self, request_id: int, queue: RoborockFuture) -> Any:
         try:
             response = await queue.async_get(self.queue_timeout)
@@ -97,9 +90,7 @@ class RoborockClient(ABC):
 
     def _async_response(self, request_id: int, protocol_id: int = 0) -> Any:
         queue = RoborockFuture(protocol_id)
-        if request_id in self._waiting_queue and not (
-            request_id == 2 and protocol_id == RoborockMessageProtocol.PING_REQUEST
-        ):
+        if request_id in self._waiting_queue:
             new_id = get_next_int(10000, 32767)
             self._logger.warning(
                 "Attempting to create a future with an existing id %s (%s)... New id is %s. "

{python_roborock-2.45.0 → python_roborock-2.47.0}/roborock/cli.py

@@ -27,6 +27,7 @@ import functools
 import json
 import logging
 import threading
+from collections.abc import Callable
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any, cast
@@ -43,6 +44,8 @@ from roborock.containers import DeviceData, HomeData, NetworkInfo, RoborockBase,
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
+from roborock.devices.traits import Trait
+from roborock.devices.traits.v1 import V1TraitMixin
 from roborock.protocol import MessageParser
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.web_api import RoborockApiClient
@@ -377,6 +380,22 @@ async def execute_scene(ctx, scene_id):
     await client.execute_scene(cache_data.user_data, scene_id)
 
 
+async def _v1_trait(context: RoborockContext, device_id: str, display_func: Callable[[], V1TraitMixin]) -> Trait:
+    device_manager = await context.get_device_manager()
+    device = await device_manager.get_device(device_id)
+    if device.v1_properties is None:
+        raise RoborockException(f"Device {device.name} does not support V1 protocol")
+
+    trait = display_func(device.v1_properties)
+    await trait.refresh()
+    return trait
+
+
+async def _display_v1_trait(context: RoborockContext, device_id: str, display_func: Callable[[], Trait]) -> None:
+    trait = await _v1_trait(context, device_id, display_func)
+    click.echo(dump_json(trait.as_dict()))
+
+
 @session.command()
 @click.option("--device_id", required=True)
 @click.pass_context
@@ -384,16 +403,7 @@ async def execute_scene(ctx, scene_id):
 async def status(ctx, device_id: str):
     """Get device status."""
     context: RoborockContext = ctx.obj
-
-    device_manager = await context.get_device_manager()
-    device = await device_manager.get_device(device_id)
-
-    if not (status_trait := device.traits.get("status")):
-        click.echo(f"Device {device.name} does not have a status trait")
-        return
-
-    status_result = await status_trait.get_status()
-    click.echo(dump_json(status_result.as_dict()))
+    await _display_v1_trait(context, device_id, lambda v1: v1.status)
 
 
 @session.command()
@@ -403,15 +413,7 @@ async def status(ctx, device_id: str):
 async def clean_summary(ctx, device_id: str):
     """Get device clean summary."""
     context: RoborockContext = ctx.obj
-
-    device_manager = await context.get_device_manager()
-    device = await device_manager.get_device(device_id)
-    if not (clean_summary_trait := device.traits.get("clean_summary")):
-        click.echo(f"Device {device.name} does not have a clean summary trait")
-        return
-
-    clean_summary_result = await clean_summary_trait.get_clean_summary()
-    click.echo(dump_json(clean_summary_result.as_dict()))
+    await _display_v1_trait(context, device_id, lambda v1: v1.clean_summary)
 
 
 @session.command()
@@ -421,17 +423,7 @@ async def clean_summary(ctx, device_id: str):
 async def volume(ctx, device_id: str):
     """Get device volume."""
     context: RoborockContext = ctx.obj
-
-    device_manager = await context.get_device_manager()
-    device = await device_manager.get_device(device_id)
-
-    if not (volume_trait := device.traits.get("sound_volume")):
-        click.echo(f"Device {device.name} does not have a volume trait")
-        return
-
-    volume_result = await volume_trait.get_volume()
-    click.echo(f"Device {device_id} volume:")
-    click.echo(volume_result)
+    await _display_v1_trait(context, device_id, lambda v1: v1.sound_volume)
 
 
 @session.command()
@@ -442,14 +434,7 @@ async def volume(ctx, device_id: str):
 async def set_volume(ctx, device_id: str, volume: int):
     """Set the devicevolume."""
     context: RoborockContext = ctx.obj
-
-    device_manager = await context.get_device_manager()
-    device = await device_manager.get_device(device_id)
-
-    if not (volume_trait := device.traits.get("sound_volume")):
-        click.echo(f"Device {device.name} does not have a volume trait")
-        return
-
+    volume_trait = await _v1_trait(context, device_id, lambda v1: v1.sound_volume)
     await volume_trait.set_volume(volume)
     click.echo(f"Set Device {device_id} volume to {volume}")
 

{python_roborock-2.45.0 → python_roborock-2.47.0}/roborock/cloud_api.py

@@ -8,6 +8,10 @@ from asyncio import Lock
 from typing import Any
 
 import paho.mqtt.client as mqtt
+from paho.mqtt.enums import MQTTErrorCode
+
+# Mypy is not seeing this for some reason. It wants me to use the depreciated ReasonCodes
+from paho.mqtt.reasoncodes import ReasonCode  # type: ignore
 
 from .api import KEEPALIVE, RoborockClient
 from .containers import DeviceData, UserData
@@ -67,7 +71,8 @@ class RoborockMqttClient(RoborockClient, ABC):
         self._mqtt_client = _Mqtt()
         self._mqtt_client.on_connect = self._mqtt_on_connect
         self._mqtt_client.on_message = self._mqtt_on_message
-        self._mqtt_client.on_disconnect = self._mqtt_on_disconnect
+        # Due to the incorrect ReasonCode, it is confused by typing
+        self._mqtt_client.on_disconnect = self._mqtt_on_disconnect  # type: ignore
         if mqtt_params.tls:
             self._mqtt_client.tls_set()
 
@@ -76,12 +81,20 @@ class RoborockMqttClient(RoborockClient, ABC):
         self._mutex = Lock()
         self._decoder: Decoder = create_mqtt_decoder(device_info.device.local_key)
         self._encoder: Encoder = create_mqtt_encoder(device_info.device.local_key)
+        self.received_message_since_last_disconnect = False
+        self._topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"
 
-    def _mqtt_on_connect(self, *args, **kwargs):
-        _, __, ___, rc, ____ = args
+    def _mqtt_on_connect(
+        self,
+        client: mqtt.Client,
+        userdata: object,
+        flags: dict[str, int],
+        rc: ReasonCode,
+        properties: mqtt.Properties | None = None,
+    ):
         connection_queue = self._waiting_queue.get(CONNECT_REQUEST_ID)
-        if rc != mqtt.MQTT_ERR_SUCCESS:
-            message = f"Failed to connect ({mqtt.error_string(rc)})"
+        if rc.is_failure:
+            message = f"Failed to connect ({rc})"
             self._logger.error(message)
             if connection_queue:
                 connection_queue.set_exception(VacuumError(message))
@@ -89,19 +102,19 @@ class RoborockMqttClient(RoborockClient, ABC):
                 self._logger.debug("Failed to notify connect future, not in queue")
             return
         self._logger.info(f"Connected to mqtt {self._mqtt_host}:{self._mqtt_port}")
-        topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"
-        (result, mid) = self._mqtt_client.subscribe(topic)
+        (result, mid) = self._mqtt_client.subscribe(self._topic)
         if result != 0:
-            message = f"Failed to subscribe ({mqtt.error_string(rc)})"
+            message = f"Failed to subscribe ({str(rc)})"
             self._logger.error(message)
             if connection_queue:
                 connection_queue.set_exception(VacuumError(message))
             return
-        self._logger.info(f"Subscribed to topic {topic}")
+        self._logger.info(f"Subscribed to topic {self._topic}")
         if connection_queue:
             connection_queue.set_result(True)
 
     def _mqtt_on_message(self, *args, **kwargs):
+        self.received_message_since_last_disconnect = True
         client, __, msg = args
         try:
             messages = self._decoder(msg.payload)
@@ -109,10 +122,16 @@ class RoborockMqttClient(RoborockClient, ABC):
         except Exception as ex:
             self._logger.exception(ex)
 
-    def _mqtt_on_disconnect(self, *args, **kwargs):
-        _, __, rc, ___ = args
+    def _mqtt_on_disconnect(
+        self,
+        client: mqtt.Client,
+        data: object,
+        flags: dict[str, int],
+        rc: ReasonCode | None,
+        properties: mqtt.Properties | None = None,
+    ):
         try:
-            exc = RoborockException(mqtt.error_string(rc)) if rc != mqtt.MQTT_ERR_SUCCESS else None
+            exc = RoborockException(str(rc)) if rc is not None and rc.is_failure else None
             super().on_connection_lost(exc)
             connection_queue = self._waiting_queue.get(DISCONNECT_REQUEST_ID)
             if connection_queue:
@@ -138,7 +157,7 @@ class RoborockMqttClient(RoborockClient, ABC):
 
         if rc != mqtt.MQTT_ERR_SUCCESS:
             disconnected_future.cancel()
-            raise RoborockException(f"Failed to disconnect ({mqtt.error_string(rc)})")
+            raise RoborockException(f"Failed to disconnect ({str(rc)})")
 
         return disconnected_future
 
@@ -178,3 +197,63 @@ class RoborockMqttClient(RoborockClient, ABC):
         )
         if info.rc != mqtt.MQTT_ERR_SUCCESS:
             raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
+
+    async def _unsubscribe(self) -> MQTTErrorCode:
+        """Unsubscribe from the topic."""
+        loop = asyncio.get_running_loop()
+        (result, mid) = await loop.run_in_executor(None, self._mqtt_client.unsubscribe, self._topic)
+
+        if result != 0:
+            message = f"Failed to unsubscribe ({mqtt.error_string(result)})"
+            self._logger.error(message)
+        else:
+            self._logger.info(f"Unsubscribed from topic {self._topic}")
+        return result
+
+    async def _subscribe(self) -> MQTTErrorCode:
+        """Subscribe to the topic."""
+        loop = asyncio.get_running_loop()
+        (result, mid) = await loop.run_in_executor(None, self._mqtt_client.subscribe, self._topic)
+
+        if result != 0:
+            message = f"Failed to subscribe ({mqtt.error_string(result)})"
+            self._logger.error(message)
+        else:
+            self._logger.info(f"Subscribed to topic {self._topic}")
+        return result
+
+    async def _reconnect(self) -> None:
+        """Reconnect to the MQTT broker."""
+        await self.async_disconnect()
+        await self.async_connect()
+
+    async def _validate_connection(self) -> None:
+        """Override the default validate connection to try to re-subscribe rather than disconnect.
+        When something seems to be wrong with our connection, we should follow the following steps:
+        1. Try to unsubscribe and resubscribe from the topic.
+        2. If we don't end up getting a message, we should completely disconnect and reconnect to the MQTT broker.
+        3. We will continue to try to disconnect and reconnect until we get a message.
+        4. If we get a message, the next time connection is lost, We will go back to step 1.
+        """
+        # If we should no longer keep the current connection alive...
+        if not self.should_keepalive():
+            self._logger.info("Resetting Roborock connection due to keepalive timeout")
+            if not self.received_message_since_last_disconnect:
+                # If we have already tried to unsub and resub, and we are still in this state,
+                # we should try to reconnect.
+                return await self._reconnect()
+            try:
+                # Mark that we have tried to unsubscribe and resubscribe
+                self.received_message_since_last_disconnect = False
+                if await self._unsubscribe() != 0:
+                    # If we fail to unsubscribe, reconnect to the broker
+                    return await self._reconnect()
+                if await self._subscribe() != 0:
+                    # If we fail to subscribe, reconnected to the broker.
+                    return await self._reconnect()
+
+            except Exception:  # noqa
+                # If we get any errors at all, we should just reconnect.
+                return await self._reconnect()
+        # Call connect to make sure everything is still in a good state.
+        await self.async_connect()

{python_roborock-2.45.0 → python_roborock-2.47.0}/roborock/devices/device.py

@@ -6,14 +6,14 @@ until the API is stable.
 
 import logging
 from abc import ABC
-from collections.abc import Callable, Mapping
-from types import MappingProxyType
+from collections.abc import Callable
 
 from roborock.containers import HomeDataDevice
 from roborock.roborock_message import RoborockMessage
 
 from .channel import Channel
-from .traits.trait import Trait
+from .traits import Trait
+from .traits.traits_mixin import TraitsMixin
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -22,19 +22,23 @@ __all__ = [
 ]
 
 
-class RoborockDevice(ABC):
+class RoborockDevice(ABC, TraitsMixin):
     """A generic channel for establishing a connection with a Roborock device.
 
     Individual channel implementations have their own methods for speaking to
     the device that hide some of the protocol specific complexity, but they
     are still specialized for the device type and protocol.
+
+    Attributes of the device are exposed through traits, which are mixed in
+    through the TraitsMixin class. Traits are optional and may not be present
+    on all devices.
     """
 
     def __init__(
         self,
         device_info: HomeDataDevice,
         channel: Channel,
-        traits: list[Trait],
+        trait: Trait,
     ) -> None:
         """Initialize the RoborockDevice.
 
@@ -42,13 +46,11 @@ class RoborockDevice(ABC):
         Use `connect()` to establish the connection, which will set up the appropriate
         protocol channel. Use `close()` to clean up all connections.
         """
+        TraitsMixin.__init__(self, trait)
         self._duid = device_info.duid
         self._name = device_info.name
         self._channel = channel
         self._unsub: Callable[[], None] | None = None
-        self._trait_map = {trait.name: trait for trait in traits}
-        if len(self._trait_map) != len(traits):
-            raise ValueError("Duplicate trait names found in traits list")
 
     @property
     def duid(self) -> str:
@@ -81,8 +83,3 @@ class RoborockDevice(ABC):
     def _on_message(self, message: RoborockMessage) -> None:
         """Handle incoming messages from the device."""
         _LOGGER.debug("Received message from device: %s", message)
-
-    @property
-    def traits(self) -> Mapping[str, Trait]:
-        """Return the traits of the device."""
-        return MappingProxyType(self._trait_map)

{python_roborock-2.45.0 → python_roborock-2.47.0}/roborock/devices/device_manager.py

@@ -7,7 +7,6 @@ from collections.abc import Awaitable, Callable
 
 import aiohttp
 
-from roborock.code_mappings import RoborockCategory
 from roborock.containers import (
     HomeData,
     HomeDataDevice,
@@ -23,14 +22,7 @@ from roborock.web_api import RoborockApiClient
 from .cache import Cache, NoCache
 from .channel import Channel
 from .mqtt_channel import create_mqtt_channel
-from .traits.b01.props import B01PropsApi
-from .traits.clean_summary import CleanSummaryTrait
-from .traits.dnd import DoNotDisturbTrait
-from .traits.dyad import DyadApi
-from .traits.sound_volume import SoundVolumeTrait
-from .traits.status import StatusTrait
-from .traits.trait import Trait
-from .traits.zeo import ZeoApi
+from .traits import Trait, a01, b01, v1
 from .v1_channel import create_v1_channel
 
 _LOGGER = logging.getLogger(__name__)
@@ -153,30 +145,20 @@ async def create_device_manager(
 
     def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
         channel: Channel
-        traits: list[Trait] = []
-        # TODO: Define a registration mechanism/factory for v1 traits
+        trait: Trait
         match device.pv:
             case DeviceVersion.V1:
                 channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
-                traits.append(StatusTrait(product, channel.rpc_channel))
-                traits.append(DoNotDisturbTrait(channel.rpc_channel))
-                traits.append(CleanSummaryTrait(channel.rpc_channel))
-                traits.append(SoundVolumeTrait(channel.rpc_channel))
+                trait = v1.create(product, channel.rpc_channel)
             case DeviceVersion.A01:
-                mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
-                match product.category:
-                    case RoborockCategory.WET_DRY_VAC:
-                        traits.append(DyadApi(mqtt_channel))
-                    case RoborockCategory.WASHING_MACHINE:
-                        traits.append(ZeoApi(mqtt_channel))
-                    case _:
-                        raise NotImplementedError(f"Device {device.name} has unsupported category {product.category}")
+                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                trait = a01.create(product, channel)
             case DeviceVersion.B01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
-                traits.append(B01PropsApi(channel))
+                trait = b01.create(channel)
             case _:
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
-        return RoborockDevice(device, channel, traits)
+        return RoborockDevice(device, channel, trait)
 
     manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session, cache=cache)
     await manager.discover_devices()

python_roborock-2.47.0/roborock/devices/traits/__init__.py

@@ -0,0 +1,15 @@
+"""Module for device traits."""
+
+from abc import ABC
+
+__all__ = [
+    "Trait",
+    "traits_mixin",
+    "v1",
+    "a01",
+    "b01",
+]
+
+
+class Trait(ABC):
+    """Base class for all traits."""

python_roborock-2.47.0/roborock/devices/traits/a01/__init__.py

@@ -0,0 +1,61 @@
+from typing import Any
+
+from roborock.containers import HomeDataProduct, RoborockCategory
+from roborock.devices.a01_channel import send_decoded_command
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.devices.traits import Trait
+from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
+
+__init__ = [
+    "DyadApi",
+    "ZeoApi",
+]
+
+
+class DyadApi(Trait):
+    """API for interacting with Dyad devices."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the Dyad API."""
+        self._channel = channel
+
+    async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
+        """Query the device for the values of the given Dyad protocols."""
+        params = {RoborockDyadDataProtocol.ID_QUERY: [int(p) for p in protocols]}
+        return await send_decoded_command(self._channel, params)
+
+    async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
+        """Set a value for a specific protocol on the device."""
+        params = {protocol: value}
+        return await send_decoded_command(self._channel, params)
+
+
+class ZeoApi(Trait):
+    """API for interacting with Zeo devices."""
+
+    name = "zeo"
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the Zeo API."""
+        self._channel = channel
+
+    async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
+        """Query the device for the values of the given protocols."""
+        params = {RoborockZeoProtocol.ID_QUERY: [int(p) for p in protocols]}
+        return await send_decoded_command(self._channel, params)
+
+    async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
+        """Set a value for a specific protocol on the device."""
+        params = {protocol: value}
+        return await send_decoded_command(self._channel, params)
+
+
+def create(product: HomeDataProduct, mqtt_channel: MqttChannel) -> DyadApi | ZeoApi:
+    """Create traits for A01 devices."""
+    match product.category:
+        case RoborockCategory.WET_DRY_VAC:
+            return DyadApi(mqtt_channel)
+        case RoborockCategory.WASHING_MACHINE:
+            return ZeoApi(mqtt_channel)
+        case _:
+            raise NotImplementedError(f"Unsupported category {product.category}")

python_roborock-2.45.0/roborock/devices/traits/b01/props.py → python_roborock-2.47.0/roborock/devices/traits/b01/__init__.py

@@ -1,26 +1,20 @@
-from __future__ import annotations
-
-import logging
+"""Traits for B01 devices."""
 
 from roborock import RoborockB01Methods
+from roborock.devices.b01_channel import send_decoded_command
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.devices.traits import Trait
 from roborock.roborock_message import RoborockB01Props
 
-from ...b01_channel import send_decoded_command
-from ...mqtt_channel import MqttChannel
-from ..trait import Trait
-
-_LOGGER = logging.getLogger(__name__)
-
-__all__ = [
-    "B01PropsApi",
+__init__ = [
+    "create_b01_traits",
+    "PropertiesApi",
 ]
 
 
-class B01PropsApi(Trait):
+class PropertiesApi(Trait):
     """API for interacting with B01 devices."""
 
-    name = "B01_props"
-
     def __init__(self, channel: MqttChannel) -> None:
         """Initialize the B01Props API."""
         self._channel = channel
@@ -30,3 +24,8 @@ class B01PropsApi(Trait):
         await send_decoded_command(
             self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params={"property": props}
         )
+
+
+def create(channel: MqttChannel) -> PropertiesApi:
+    """Create traits for B01 devices."""
+    return PropertiesApi(channel)

python_roborock-2.47.0/roborock/devices/traits/traits_mixin.py

@@ -0,0 +1,61 @@
+"""Holds device traits mixin and related code.
+
+This holds the TraitsMixin class, which is used to provide accessors for
+various device traits. Each trait is a class that encapsulates a specific
+set of functionality for a device, such as controlling a vacuum or a mop.
+
+The TraitsMixin holds traits across all protocol types. A trait is supported
+if it is non-None.
+"""
+
+from dataclasses import dataclass, fields
+from typing import get_args, get_origin
+
+from . import Trait, a01, b01, v1
+
+__all__ = [
+    "TraitsMixin",
+]
+
+
+@dataclass(init=False)
+class TraitsMixin:
+    """Mixin to provide trait accessors."""
+
+    v1_properties: v1.PropertiesApi | None = None
+    """V1 properties trait, if supported."""
+
+    dyad: a01.DyadApi | None = None
+    """Dyad API, if supported."""
+
+    zeo: a01.ZeoApi | None = None
+    """Zeo API, if supported."""
+
+    b01_properties: b01.PropertiesApi | None = None
+    """B01 properties trait, if supported."""
+
+    def __init__(self, trait: Trait) -> None:
+        """Initialize the TraitsMixin with the given trait.
+
+        This will populate the appropriate trait attributes based on the types
+        of the traits provided.
+        """
+        for item in fields(self):
+            trait_type = _get_trait_type(item)
+            if trait_type == type(trait):
+                setattr(self, item.name, trait)
+                break
+
+
+def _get_trait_type(item) -> type[Trait]:
+    """Get the trait type from a dataclass field."""
+    if get_origin(item.type) is None:
+        raise ValueError(f"Trait {item.name} is not an optional type")
+    if (args := get_args(item.type)) is None:
+        raise ValueError(f"Trait {item.name} is not an optional type")
+    if len(args) != 2 or args[1] is not type(None):
+        raise ValueError(f"Trait {item.name} is not an optional type")
+    trait_type = args[0]
+    if not issubclass(trait_type, Trait):
+        raise ValueError(f"Trait {item.name} is not a Trait subclass")
+    return trait_type