Small cleanup of Shodan (#64467)

homeassistant/components/shodan/sensor.py

@@ -8,7 +8,7 @@ import shodan
 import voluptuous as vol
 
 from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
-from homeassistant.const import ATTR_ATTRIBUTION, CONF_API_KEY, CONF_NAME
+from homeassistant.const import CONF_API_KEY, CONF_NAME
 from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -16,14 +16,10 @@ from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 _LOGGER = logging.getLogger(__name__)
 
-ATTRIBUTION = "Data provided by Shodan"
-
 CONF_QUERY = "query"
 
 DEFAULT_NAME = "Shodan Sensor"
 
-ICON = "mdi:tooltip-text"
-
 SCAN_INTERVAL = timedelta(minutes=15)
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
@@ -42,9 +38,9 @@ def setup_platform(
     discovery_info: DiscoveryInfoType | None = None,
 ) -> None:
     """Set up the Shodan sensor."""
-    api_key = config.get(CONF_API_KEY)
-    name = config.get(CONF_NAME)
-    query = config.get(CONF_QUERY)
+    api_key = config[CONF_API_KEY]
+    name = config[CONF_NAME]
+    query = config[CONF_QUERY]
 
     data = ShodanData(shodan.Shodan(api_key), query)
     try:
@@ -59,53 +55,29 @@ def setup_platform(
 class ShodanSensor(SensorEntity):
     """Representation of the Shodan sensor."""
 
-    def __init__(self, data, name):
+    _attr_attribution = "Data provided by Shodan"
+    _attr_icon = "mdi:tooltip-text"
+    _attr_native_unit_of_measurement = "Hits"
+
+    def __init__(self, data: ShodanData, name: str) -> None:
         """Initialize the Shodan sensor."""
         self.data = data
-        self._name = name
-        self._state = None
-        self._unit_of_measurement = "Hits"
+        self._attr_name = name
 
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._state
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit the value is expressed in."""
-        return self._unit_of_measurement
-
-    @property
-    def icon(self):
-        """Return the icon to use in the frontend, if any."""
-        return ICON
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes of the sensor."""
-        return {ATTR_ATTRIBUTION: ATTRIBUTION}
-
-    def update(self):
+    def update(self) -> None:
         """Get the latest data and updates the states."""
-        self.data.update()
-        self._state = self.data.details["total"]
+        data = self.data.update()
+        self._attr_native_value = data.details["total"]
 
 
 class ShodanData:
     """Get the latest data and update the states."""
 
-    def __init__(self, api, query):
+    def __init__(self, api: shodan.Shodan, query: str) -> None:
         """Initialize the data object."""
         self._api = api
         self._query = query
-        self.details = None
 
     def update(self):
         """Get the latest data from shodan.io."""
-        self.details = self._api.count(self._query)
+        return self._api.count(self._query)