94 lines
		
	
	
		
			3.0 KiB
		
	
	
	
		
			Python
		
	
	
			
		
		
	
	
			94 lines
		
	
	
		
			3.0 KiB
		
	
	
	
		
			Python
		
	
	
"""Support for sending data to StatsD."""
 | 
						|
import logging
 | 
						|
 | 
						|
import statsd
 | 
						|
import voluptuous as vol
 | 
						|
 | 
						|
from homeassistant.const import CONF_HOST, CONF_PORT, CONF_PREFIX, EVENT_STATE_CHANGED
 | 
						|
from homeassistant.core import HomeAssistant
 | 
						|
from homeassistant.helpers import state as state_helper
 | 
						|
import homeassistant.helpers.config_validation as cv
 | 
						|
from homeassistant.helpers.typing import ConfigType
 | 
						|
 | 
						|
_LOGGER = logging.getLogger(__name__)
 | 
						|
 | 
						|
CONF_ATTR = "log_attributes"
 | 
						|
CONF_RATE = "rate"
 | 
						|
CONF_VALUE_MAP = "value_mapping"
 | 
						|
 | 
						|
DEFAULT_HOST = "localhost"
 | 
						|
DEFAULT_PORT = 8125
 | 
						|
DEFAULT_PREFIX = "hass"
 | 
						|
DEFAULT_RATE = 1
 | 
						|
DOMAIN = "statsd"
 | 
						|
 | 
						|
CONFIG_SCHEMA = vol.Schema(
 | 
						|
    {
 | 
						|
        DOMAIN: vol.Schema(
 | 
						|
            {
 | 
						|
                vol.Required(CONF_HOST, default=DEFAULT_HOST): cv.string,
 | 
						|
                vol.Optional(CONF_ATTR, default=False): cv.boolean,
 | 
						|
                vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
 | 
						|
                vol.Optional(CONF_PREFIX, default=DEFAULT_PREFIX): cv.string,
 | 
						|
                vol.Optional(CONF_RATE, default=DEFAULT_RATE): vol.All(
 | 
						|
                    vol.Coerce(int), vol.Range(min=1)
 | 
						|
                ),
 | 
						|
                vol.Optional(CONF_VALUE_MAP): dict,
 | 
						|
            }
 | 
						|
        )
 | 
						|
    },
 | 
						|
    extra=vol.ALLOW_EXTRA,
 | 
						|
)
 | 
						|
 | 
						|
 | 
						|
def setup(hass: HomeAssistant, config: ConfigType) -> bool:
 | 
						|
    """Set up the StatsD component."""
 | 
						|
 | 
						|
    conf = config[DOMAIN]
 | 
						|
    host = conf.get(CONF_HOST)
 | 
						|
    port = conf.get(CONF_PORT)
 | 
						|
    sample_rate = conf.get(CONF_RATE)
 | 
						|
    prefix = conf.get(CONF_PREFIX)
 | 
						|
    value_mapping = conf.get(CONF_VALUE_MAP)
 | 
						|
    show_attribute_flag = conf.get(CONF_ATTR)
 | 
						|
 | 
						|
    statsd_client = statsd.StatsClient(host=host, port=port, prefix=prefix)
 | 
						|
 | 
						|
    def statsd_event_listener(event):
 | 
						|
        """Listen for new messages on the bus and sends them to StatsD."""
 | 
						|
        if (state := event.data.get("new_state")) is None:
 | 
						|
            return
 | 
						|
 | 
						|
        try:
 | 
						|
            if value_mapping and state.state in value_mapping:
 | 
						|
                _state = float(value_mapping[state.state])
 | 
						|
            else:
 | 
						|
                _state = state_helper.state_as_number(state)
 | 
						|
        except ValueError:
 | 
						|
            # Set the state to none and continue for any numeric attributes.
 | 
						|
            _state = None
 | 
						|
 | 
						|
        states = dict(state.attributes)
 | 
						|
 | 
						|
        _LOGGER.debug("Sending %s", state.entity_id)
 | 
						|
 | 
						|
        if show_attribute_flag is True:
 | 
						|
            if isinstance(_state, (float, int)):
 | 
						|
                statsd_client.gauge(f"{state.entity_id}.state", _state, sample_rate)
 | 
						|
 | 
						|
            # Send attribute values
 | 
						|
            for key, value in states.items():
 | 
						|
                if isinstance(value, (float, int)):
 | 
						|
                    stat = "{}.{}".format(state.entity_id, key.replace(" ", "_"))
 | 
						|
                    statsd_client.gauge(stat, value, sample_rate)
 | 
						|
 | 
						|
        elif isinstance(_state, (float, int)):
 | 
						|
            statsd_client.gauge(state.entity_id, _state, sample_rate)
 | 
						|
 | 
						|
        # Increment the count
 | 
						|
        statsd_client.incr(state.entity_id, rate=sample_rate)
 | 
						|
 | 
						|
    hass.bus.listen(EVENT_STATE_CHANGED, statsd_event_listener)
 | 
						|
 | 
						|
    return True
 |