Cleanup and reduce duplicate code from recent template changes (#40012)

As a result of refactoring, there is duplicate code
we can now reduce.

Additionally `_wrap_state` can be removed because
it had unreachable checks for `None`
This commit is contained in:
J. Nick Koston 2020-09-13 09:21:11 -05:00 committed by GitHub
parent 84578f515d
commit ff4bb962c4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -6,6 +6,7 @@ from functools import wraps
import json import json
import logging import logging
import math import math
from operator import attrgetter
import random import random
import re import re
from typing import Any, Iterable, List, Optional, Union from typing import Any, Iterable, List, Optional, Union
@ -424,12 +425,7 @@ class AllStates:
def __iter__(self): def __iter__(self):
"""Return all states.""" """Return all states."""
self._collect_all() self._collect_all()
return iter( return _state_iterator(self._hass, None)
_wrap_state(self._hass, state)
for state in sorted(
self._hass.states.async_all(), key=lambda state: state.entity_id
)
)
def __len__(self) -> int: def __len__(self) -> int:
"""Return number of states.""" """Return number of states."""
@ -469,15 +465,7 @@ class DomainStates:
def __iter__(self): def __iter__(self):
"""Return the iteration over all the states.""" """Return the iteration over all the states."""
self._collect_domain() self._collect_domain()
return iter( return _state_iterator(self._hass, self._domain)
sorted(
(
_wrap_state(self._hass, state)
for state in self._hass.states.async_all(self._domain)
),
key=lambda state: state.entity_id,
)
)
def __len__(self) -> int: def __len__(self) -> int:
"""Return number of states.""" """Return number of states."""
@ -492,6 +480,8 @@ class DomainStates:
class TemplateState(State): class TemplateState(State):
"""Class to represent a state object in a template.""" """Class to represent a state object in a template."""
__slots__ = ("_hass", "_state")
# Inheritance is done so functions that check against State keep working # Inheritance is done so functions that check against State keep working
# pylint: disable=super-init-not-called # pylint: disable=super-init-not-called
def __init__(self, hass, state): def __init__(self, hass, state):
@ -547,11 +537,12 @@ def _collect_state(hass: HomeAssistantType, entity_id: str) -> None:
entity_collect.entities.add(entity_id) entity_collect.entities.add(entity_id)
def _wrap_state( def _state_iterator(hass: HomeAssistantType, domain: Optional[str]) -> Iterable:
hass: HomeAssistantType, state: Optional[State] """Create an state iterator for a domain or all states."""
) -> Optional[TemplateState]: return iter(
"""Wrap a state.""" TemplateState(hass, state)
return None if state is None else TemplateState(hass, state) for state in sorted(hass.states.async_all(domain), key=attrgetter("entity_id"))
)
def _get_state(hass: HomeAssistantType, entity_id: str) -> Optional[TemplateState]: def _get_state(hass: HomeAssistantType, entity_id: str) -> Optional[TemplateState]:
@ -561,7 +552,7 @@ def _get_state(hass: HomeAssistantType, entity_id: str) -> Optional[TemplateStat
# access to the state properties in the state wrapper. # access to the state properties in the state wrapper.
_collect_state(hass, entity_id) _collect_state(hass, entity_id)
return None return None
return _wrap_state(hass, state) return TemplateState(hass, state)
def _resolve_state( def _resolve_state(