Google Assistant Trait Suggestions

Just recently I have updated the google_assistant/trait.py code to handle some specific and common vacuum commands (mine being the roborock S6) to support Locator, StartStop (zones) and EnergyStorage. With these traits being added/modified I am now able to query the battery level, clean by zone/segments and locate my vacuum cleaner which pretty much surpasses even the Xiaomi Home app yet very simple to integrate.

Furrther more, I found overriding the google_assistant/trait.py file to be bad practice as this can be prone to problems in future updates etc. So could it be viable to have some kind of way where the integration checks to see if another integration has an override trait? Simply put, this could be a simple check to see if an integration has a method for a certain trait in which it can call instead of resorting to the default. I say this because I know you can already add traits with @register_trait which is fine but for a trait already defined like StartStopTrait would create the same trait twice if an integration wanted to have its own customized version of StartStopTrait. Maybe even just removing duplicates in the GoogleEntity _traits list?

I understand the universibility behind a constant set of traits but limits those integrations that has extra features from its full potential when it comes to Google Assistant and the commands it could give.

Below is my added traits which should be self explainatory but would like the moderators/coders of google_assistant to consider adding these for additional functionality into the Google Assistant eco.

TRAIT_LOCATOR = f"{PREFIX_TRAITS}Locator"
TRAIT_ENERGY_STORAGE = f"{PREFIX_TRAITS}EnergyStorage"

COMMAND_LOCATE = f"{PREFIX_COMMANDS}Locate"
COMMAND_CHARGE = f"{PREFIX_COMMANDS}Charge"


@register_trait
class VacuumEnergyStorageTrait(_Trait):

    name = TRAIT_ENERGY_STORAGE
    commands = [COMMAND_CHARGE]

    @staticmethod
    def supported(domain, features, device_class):
        if domain == vacuum.DOMAIN:
            return features & vacuum.SUPPORT_BATTERY
        return False

    def sync_attributes(self):
        return {
            "isRechargeable": False,
            "queryOnlyEnergyStorage": True
        }

    def query_attributes(self):
        attrs = self.state.attributes
        return {
            'capacityRemaining': [{
                "rawValue": attrs.get('battery_level'),
                "unit": "PERCENTAGE"
            }],
            'isCharging': attrs.get('status') == 'Charging',
            'isPluggedIn': self.state.state == vacuum.STATE_DOCKED
        }

    async def execute(self, command, data, params, challenge):
        _LOGGER.error('Command Charge Executed: %s', command)

@register_trait
class VacuumFanSpeedTrait(_Trait):

    name = TRAIT_FANSPEED
    commands = [COMMAND_FANSPEED]

    @staticmethod
    def supported(domain, features, device_class):
        if domain == vacuum.DOMAIN:
            return features & vacuum.SUPPORT_FAN_SPEED
        return False

    def sync_attributes(self):
        speeds = []
        modes = self.state.attributes.get(vacuum.ATTR_FAN_SPEED_LIST, [])
        for mode in modes:
            speed = {
                "speed_name": mode,
                "speed_values": [{"speed_synonym": [mode], "lang": "en"}],
            }
            speeds.append(speed)

        return {
            "availableFanSpeeds": {"speeds": speeds, "ordered": True},
            "reversible": False,
        }

    def query_attributes(self):
        speed = self.state.attributes.get(vacuum.ATTR_FAN_SPEED)
        response = {}
        if speed is not None:
            response["currentFanSpeedSetting"] = speed
        return response

    async def execute(self, command, data, params, challenge):
        domain = self.state.domain
        if domain == vacuum.DOMAIN:
            await self.hass.services.async_call(
                vacuum.DOMAIN,
                vacuum.SERVICE_SET_FAN_SPEED,
                {
                    ATTR_ENTITY_ID: self.state.entity_id,
                    vacuum.ATTR_FAN_SPEED: params["fanSpeed"],
                },
                blocking=True,
                context=data.context,
            )

@register_trait
class VacuumLocatorTrait(_Trait):

    name = TRAIT_LOCATOR
    commands = [COMMAND_LOCATE]

    @staticmethod
    def supported(domain, features, device_class):
        if domain == vacuum.DOMAIN:
            return features & vacuum.SUPPORT_LOCATE
        return False

    def sync_attributes(self):
        return {}

    def query_attributes(self):
        return {}

    async def execute(self, command, data, params, challenge):
        await self.hass.services.async_call(
            vacuum.DOMAIN,
            vacuum.SERVICE_LOCATE,
            {ATTR_ENTITY_ID: self.state.entity_id},
            blocking=True,
            context=data.context,
        )

Additionally I modified the StartStopTrait to support zone cleaning with the vacuum. I’ve created another post in hopes that the Vacuum Entity will support zones so that this can be integrated in the future and allow zone cleaning through Google Assistant natively. https://community.home-assistant.io/t/vacuum-entity-to-support-zones/259904

VACUUM_ROOMS = {
    'hallway': 21,
    'family room': 27,
    'living room': 26,
    'bathroom': 19,
    'toilet': 18,
    'kitchen': 16,
    'study': 17,
    'master bedroom': 24,
    'kids bedroom': 20,
    'laundry': 23,
    'entrance': 25
}
    @staticmethod
    def supported(domain, features, device_class):
        """Test if state is supported."""
        if domain == vacuum.DOMAIN:
            return True

        if domain == cover.DOMAIN and features & cover.SUPPORT_STOP:
            return True

        return False

    def sync_attributes(self):
        """Return StartStop attributes for a sync request."""
        domain = self.state.domain
        if domain == vacuum.DOMAIN:
            return {
                "pausable": self.state.attributes.get(ATTR_SUPPORTED_FEATURES, 0) & vacuum.SUPPORT_PAUSE != 0,
                "availableZones": list(VACUUM_ROOMS.keys())
            }
        if domain == cover.DOMAIN:
            return {}

    def query_attributes(self):
        """Return StartStop query attributes."""
        domain = self.state.domain
        state = self.state.state

        if domain == vacuum.DOMAIN:
            return {
                "isRunning": state == vacuum.STATE_CLEANING,
                "isPaused": state == vacuum.STATE_PAUSED,
            }

        if domain == cover.DOMAIN:
            return {"isRunning": state in (cover.STATE_CLOSING, cover.STATE_OPENING)}

    async def execute(self, command, data, params, challenge):
        """Execute a StartStop command."""
        domain = self.state.domain
        if domain == vacuum.DOMAIN:
            return await self._execute_vacuum(command, data, params, challenge)
        if domain == cover.DOMAIN:
            return await self._execute_cover(command, data, params, challenge)

    async def _execute_vacuum(self, command, data, params, challenge):
        """Execute a StartStop command."""
        if command == COMMAND_STARTSTOP:
            if params["start"]:
                if "zone" in params or "multipleZones" in params:
                    zones = []
                    if "multipleZones" in params:
                        for zone in params["multipleZones"]:
                            zones.append( VACUUM_ROOMS.get(zone) )
                    else:
                        zones.append( VACUUM_ROOMS.get(params["zone"]) )

                    await self.hass.services.async_call(
                        'xiaomi_miio',
                        'vacuum_clean_segment',
                        {
                            ATTR_ENTITY_ID: self.state.entity_id,
                            'segments': zones
                        },
                        blocking=True,
                        context=data.context,
                    )
                else:
                    await self.hass.services.async_call(
                        self.state.domain,
                        vacuum.SERVICE_START,
                        {ATTR_ENTITY_ID: self.state.entity_id},
                        blocking=True,
                        context=data.context,
                    )
            else:
                await self.hass.services.async_call(
                    self.state.domain,
                    vacuum.SERVICE_STOP,
                    {ATTR_ENTITY_ID: self.state.entity_id},
                    blocking=True,
                    context=data.context,
                )
        elif command == COMMAND_PAUSEUNPAUSE:
            if params["pause"]:
                await self.hass.services.async_call(
                    self.state.domain,
                    vacuum.SERVICE_PAUSE,
                    {ATTR_ENTITY_ID: self.state.entity_id},
                    blocking=True,
                    context=data.context,
                )
            else:
                await self.hass.services.async_call(
                    self.state.domain,
                    vacuum.SERVICE_START,
                    {ATTR_ENTITY_ID: self.state.entity_id},
                    blocking=True,
                    context=data.context,
                )

    async def _execute_cover(self, command, data, params, challenge):
        """Execute a StartStop command."""
        if command == COMMAND_STARTSTOP:
            if params["start"] is False:
                if self.state.state in (cover.STATE_CLOSING, cover.STATE_OPENING):
                    await self.hass.services.async_call(
                        self.state.domain,
                        cover.SERVICE_STOP_COVER,
                        {ATTR_ENTITY_ID: self.state.entity_id},
                        blocking=True,
                        context=data.context,
                    )
                else:
                    raise SmartHomeError(
                        ERR_ALREADY_STOPPED, "Cover is already stopped"
                    )
            else:
                raise SmartHomeError(
                    ERR_NOT_SUPPORTED, "Starting a cover is not supported"
                )
        else:
            raise SmartHomeError(
                ERR_NOT_SUPPORTED, f"Command {command} is not supported"
            )