From 733f1e1101b2f51adaef7cee557d4cd12171885a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ville=20Skytt=C3=A4?= <ville.skytta@iki.fi>
Date: Sat, 15 Feb 2020 23:03:53 +0200
Subject: [PATCH] Helpers typing improvements (#31865)

---
 homeassistant/helpers/restore_state.py | 18 +++++++++---------
 homeassistant/helpers/template.py      | 10 +++++-----
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/homeassistant/helpers/restore_state.py b/homeassistant/helpers/restore_state.py
index 7b2ec8c7f75..d57d3ad9920 100644
--- a/homeassistant/helpers/restore_state.py
+++ b/homeassistant/helpers/restore_state.py
@@ -2,7 +2,7 @@
 import asyncio
 from datetime import datetime, timedelta
 import logging
-from typing import Any, Dict, List, Optional, Set
+from typing import Any, Awaitable, Dict, List, Optional, Set, cast
 
 from homeassistant.const import EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import (
@@ -20,9 +20,6 @@ from homeassistant.helpers.json import JSONEncoder
 from homeassistant.helpers.storage import Store
 import homeassistant.util.dt as dt_util
 
-# mypy: allow-untyped-calls, allow-untyped-defs, no-check-untyped-defs
-# mypy: no-warn-return-any
-
 DATA_RESTORE_STATE_TASK = "restore_state_task"
 
 _LOGGER = logging.getLogger(__name__)
@@ -45,7 +42,7 @@ class StoredState:
         self.state = state
         self.last_seen = last_seen
 
-    def as_dict(self) -> Dict:
+    def as_dict(self) -> Dict[str, Any]:
         """Return a dict representation of the stored state."""
         return {"state": self.state.as_dict(), "last_seen": self.last_seen}
 
@@ -104,7 +101,7 @@ class RestoreStateData:
                 load_instance(hass)
             )
 
-        return await task
+        return await cast(Awaitable["RestoreStateData"], task)
 
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the restore state data class."""
@@ -211,15 +208,18 @@ class RestoreStateData:
         self.entity_ids.remove(entity_id)
 
 
-def _encode(value):
+def _encode(value: Any) -> Any:
     """Little helper to JSON encode a value."""
     try:
-        return JSONEncoder.default(None, value)
+        return JSONEncoder.default(
+            None,  # type: ignore
+            value,
+        )
     except TypeError:
         return value
 
 
-def _encode_complex(value):
+def _encode_complex(value: Any) -> Any:
     """Recursively encode all values with the JSONEncoder."""
     if isinstance(value, dict):
         return {_encode(key): _encode_complex(value) for key, value in value.items()}
diff --git a/homeassistant/helpers/template.py b/homeassistant/helpers/template.py
index 77642ce5052..e7f89b482e2 100644
--- a/homeassistant/helpers/template.py
+++ b/homeassistant/helpers/template.py
@@ -378,7 +378,7 @@ class DomainStates:
             raise TemplateError(f"Invalid entity ID '{entity_id}'")
         return _get_state(self._hass, entity_id)
 
-    def _collect_domain(self):
+    def _collect_domain(self) -> None:
         entity_collect = self._hass.data.get(_RENDER_INFO)
         if entity_collect is not None:
             # pylint: disable=protected-access
@@ -398,12 +398,12 @@ class DomainStates:
             )
         )
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Return number of states."""
         self._collect_domain()
         return len(self._hass.states.async_entity_ids(self._domain))
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Representation of Domain States."""
         return f"<template DomainStates('{self._domain}')>"
 
@@ -426,7 +426,7 @@ class TemplateState(State):
         return state
 
     @property
-    def state_with_unit(self):
+    def state_with_unit(self) -> str:
         """Return the state concatenated with the unit if available."""
         state = object.__getattribute__(self, "_access_state")()
         unit = state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)
@@ -447,7 +447,7 @@ class TemplateState(State):
         state = object.__getattribute__(self, "_access_state")()
         return getattr(state, name)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Representation of Template State."""
         state = object.__getattribute__(self, "_access_state")()
         rep = state.__repr__()