From 6f9ccb54342fdd723c769724da0f1e9ea98894a1 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <frenck@frenck.nl>
Date: Tue, 24 Sep 2019 23:20:04 +0200
Subject: [PATCH] Add and corrects typehints in Entity helper & core class
 (#26805)

* Add and corrects typehints in Entity class

* Adjust state type based on comments
---
 homeassistant/core.py           |  8 ++++----
 homeassistant/helpers/entity.py | 28 ++++++++++++++--------------
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/homeassistant/core.py b/homeassistant/core.py
index c29d41ace9a..31761f2560f 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -703,7 +703,7 @@ class State:
     def __init__(
         self,
         entity_id: str,
-        state: Any,
+        state: str,
         attributes: Optional[Dict] = None,
         last_changed: Optional[datetime.datetime] = None,
         last_updated: Optional[datetime.datetime] = None,
@@ -732,7 +732,7 @@ class State:
             )
 
         self.entity_id = entity_id.lower()
-        self.state: str = state
+        self.state = state
         self.attributes = MappingProxyType(attributes or {})
         self.last_updated = last_updated or dt_util.utcnow()
         self.last_changed = last_changed or self.last_updated
@@ -924,7 +924,7 @@ class StateMachine:
     def set(
         self,
         entity_id: str,
-        new_state: Any,
+        new_state: str,
         attributes: Optional[Dict] = None,
         force_update: bool = False,
         context: Optional[Context] = None,
@@ -950,7 +950,7 @@ class StateMachine:
     def async_set(
         self,
         entity_id: str,
-        new_state: Any,
+        new_state: str,
         attributes: Optional[Dict] = None,
         force_update: bool = False,
         context: Optional[Context] = None,
diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py
index af8d5589c8a..4911c5d5fb9 100644
--- a/homeassistant/helpers/entity.py
+++ b/homeassistant/helpers/entity.py
@@ -3,7 +3,7 @@ from datetime import timedelta
 import logging
 import functools as ft
 from timeit import default_timer as timer
-from typing import Any, Optional, List, Iterable
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 from homeassistant.const import (
     ATTR_ASSUMED_STATE,
@@ -26,7 +26,7 @@ from homeassistant.helpers.entity_registry import (
     EVENT_ENTITY_REGISTRY_UPDATED,
     RegistryEntry,
 )
-from homeassistant.core import HomeAssistant, callback, CALLBACK_TYPE
+from homeassistant.core import HomeAssistant, callback, CALLBACK_TYPE, Context
 from homeassistant.config import DATA_CUSTOMIZE
 from homeassistant.exceptions import NoEntitySpecifiedError
 from homeassistant.util import ensure_unique_string, slugify
@@ -137,12 +137,12 @@ class Entity:
         return None
 
     @property
-    def state(self) -> str:
+    def state(self) -> Union[None, str, int, float]:
         """Return the state of the entity."""
         return STATE_UNKNOWN
 
     @property
-    def state_attributes(self):
+    def state_attributes(self) -> Optional[Dict[str, Any]]:
         """Return the state attributes.
 
         Implemented by component base class.
@@ -150,7 +150,7 @@ class Entity:
         return None
 
     @property
-    def device_state_attributes(self):
+    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
         """Return device specific state attributes.
 
         Implemented by platform classes.
@@ -158,7 +158,7 @@ class Entity:
         return None
 
     @property
-    def device_info(self):
+    def device_info(self) -> Optional[Dict[str, Any]]:
         """Return device specific attributes.
 
         Implemented by platform classes.
@@ -171,17 +171,17 @@ class Entity:
         return None
 
     @property
-    def unit_of_measurement(self):
+    def unit_of_measurement(self) -> Optional[str]:
         """Return the unit of measurement of this entity, if any."""
         return None
 
     @property
-    def icon(self):
+    def icon(self) -> Optional[str]:
         """Return the icon to use in the frontend, if any."""
         return None
 
     @property
-    def entity_picture(self):
+    def entity_picture(self) -> Optional[str]:
         """Return the entity picture to use in the frontend, if any."""
         return None
 
@@ -215,12 +215,12 @@ class Entity:
         return None
 
     @property
-    def context_recent_time(self):
+    def context_recent_time(self) -> timedelta:
         """Time that a context is considered recent."""
         return timedelta(seconds=5)
 
     @property
-    def entity_registry_enabled_default(self):
+    def entity_registry_enabled_default(self) -> bool:
         """Return if the entity should be enabled when first added to the entity registry."""
         return True
 
@@ -230,12 +230,12 @@ class Entity:
     # produce undesirable effects in the entity's operation.
 
     @property
-    def enabled(self):
+    def enabled(self) -> bool:
         """Return if the entity is enabled in the entity registry."""
         return self.registry_entry is None or not self.registry_entry.disabled
 
     @callback
-    def async_set_context(self, context):
+    def async_set_context(self, context: Context) -> None:
         """Set the context the entity currently operates under."""
         self._context = context
         self._context_set = dt_util.utcnow()
@@ -540,7 +540,7 @@ class Entity:
 
         return self.unique_id == other.unique_id
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return the representation."""
         return "<Entity {}: {}>".format(self.name, self.state)