From abf6edea6aaf5bdeb3205240bbf8c790e8058810 Mon Sep 17 00:00:00 2001
From: Lukas <lukas@hetzenecker.me>
Date: Tue, 26 Oct 2021 13:39:34 +0200
Subject: [PATCH] Add device_class and state_class as optional attributes to
 the scrape sensor, to support statistics (#58164)

---
 homeassistant/components/scrape/sensor.py | 57 +++++++++++++++++------
 1 file changed, 42 insertions(+), 15 deletions(-)

diff --git a/homeassistant/components/scrape/sensor.py b/homeassistant/components/scrape/sensor.py
index 2e118bf744f..330df3dd245 100644
--- a/homeassistant/components/scrape/sensor.py
+++ b/homeassistant/components/scrape/sensor.py
@@ -6,9 +6,16 @@ import httpx
 import voluptuous as vol
 
 from homeassistant.components.rest.data import RestData
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    DEVICE_CLASSES_SCHEMA,
+    PLATFORM_SCHEMA,
+    STATE_CLASSES_SCHEMA,
+    SensorEntity,
+)
 from homeassistant.const import (
     CONF_AUTHENTICATION,
+    CONF_DEVICE_CLASS,
     CONF_HEADERS,
     CONF_NAME,
     CONF_PASSWORD,
@@ -45,6 +52,8 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
         vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
         vol.Optional(CONF_PASSWORD): cv.string,
         vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
+        vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
+        vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA,
         vol.Optional(CONF_USERNAME): cv.string,
         vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
         vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): cv.boolean,
@@ -64,6 +73,8 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
     attr = config.get(CONF_ATTR)
     index = config.get(CONF_INDEX)
     unit = config.get(CONF_UNIT_OF_MEASUREMENT)
+    device_class = config.get(CONF_DEVICE_CLASS)
+    state_class = config.get(CONF_STATE_CLASS)
     username = config.get(CONF_USERNAME)
     password = config.get(CONF_PASSWORD)
 
@@ -84,33 +95,49 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         raise PlatformNotReady
 
     async_add_entities(
-        [ScrapeSensor(rest, name, select, attr, index, value_template, unit)], True
+        [
+            ScrapeSensor(
+                rest,
+                name,
+                select,
+                attr,
+                index,
+                value_template,
+                unit,
+                device_class,
+                state_class,
+            )
+        ],
+        True,
     )
 
 
 class ScrapeSensor(SensorEntity):
     """Representation of a web scrape sensor."""
 
-    def __init__(self, rest, name, select, attr, index, value_template, unit):
+    def __init__(
+        self,
+        rest,
+        name,
+        select,
+        attr,
+        index,
+        value_template,
+        unit,
+        device_class,
+        state_class,
+    ):
         """Initialize a web scrape sensor."""
         self.rest = rest
-        self._name = name
         self._state = None
         self._select = select
         self._attr = attr
         self._index = index
         self._value_template = value_template
-        self._unit_of_measurement = unit
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
-    @property
-    def native_unit_of_measurement(self):
-        """Return the unit the value is expressed in."""
-        return self._unit_of_measurement
+        self._attr_name = name
+        self._attr_native_unit_of_measurement = unit
+        self._attr_device_class = device_class
+        self._attr_state_class = state_class
 
     @property
     def native_value(self):