[core] Add typing to some core files (#10843)

esphome/automation.py

@@ -15,7 +15,10 @@ from esphome.const import (
     CONF_TYPE_ID,
     CONF_UPDATE_INTERVAL,
 )
+from esphome.core import ID
+from esphome.cpp_generator import MockObj, MockObjClass, TemplateArgsType
 from esphome.schema_extractors import SCHEMA_EXTRACT, schema_extractor
+from esphome.types import ConfigType
 from esphome.util import Registry
 
 
@@ -49,11 +52,11 @@ def maybe_conf(conf, *validators):
     return validate
 
 
-def register_action(name, action_type, schema):
+def register_action(name: str, action_type: MockObjClass, schema: cv.Schema):
     return ACTION_REGISTRY.register(name, action_type, schema)
 
 
-def register_condition(name, condition_type, schema):
+def register_condition(name: str, condition_type: MockObjClass, schema: cv.Schema):
     return CONDITION_REGISTRY.register(name, condition_type, schema)
 
 
@@ -164,43 +167,78 @@ XorCondition = cg.esphome_ns.class_("XorCondition", Condition)
 
 
 @register_condition("and", AndCondition, validate_condition_list)
-async def and_condition_to_code(config, condition_id, template_arg, args):
+async def and_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     conditions = await build_condition_list(config, template_arg, args)
     return cg.new_Pvariable(condition_id, template_arg, conditions)
 
 
 @register_condition("or", OrCondition, validate_condition_list)
-async def or_condition_to_code(config, condition_id, template_arg, args):
+async def or_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     conditions = await build_condition_list(config, template_arg, args)
     return cg.new_Pvariable(condition_id, template_arg, conditions)
 
 
 @register_condition("all", AndCondition, validate_condition_list)
-async def all_condition_to_code(config, condition_id, template_arg, args):
+async def all_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     conditions = await build_condition_list(config, template_arg, args)
     return cg.new_Pvariable(condition_id, template_arg, conditions)
 
 
 @register_condition("any", OrCondition, validate_condition_list)
-async def any_condition_to_code(config, condition_id, template_arg, args):
+async def any_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     conditions = await build_condition_list(config, template_arg, args)
     return cg.new_Pvariable(condition_id, template_arg, conditions)
 
 
 @register_condition("not", NotCondition, validate_potentially_and_condition)
-async def not_condition_to_code(config, condition_id, template_arg, args):
+async def not_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     condition = await build_condition(config, template_arg, args)
     return cg.new_Pvariable(condition_id, template_arg, condition)
 
 
 @register_condition("xor", XorCondition, validate_condition_list)
-async def xor_condition_to_code(config, condition_id, template_arg, args):
+async def xor_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     conditions = await build_condition_list(config, template_arg, args)
     return cg.new_Pvariable(condition_id, template_arg, conditions)
 
 
 @register_condition("lambda", LambdaCondition, cv.returning_lambda)
-async def lambda_condition_to_code(config, condition_id, template_arg, args):
+async def lambda_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     lambda_ = await cg.process_lambda(config, args, return_type=bool)
     return cg.new_Pvariable(condition_id, template_arg, lambda_)
 
@@ -217,7 +255,12 @@ async def lambda_condition_to_code(config, condition_id, template_arg, args):
         }
     ).extend(cv.COMPONENT_SCHEMA),
 )
-async def for_condition_to_code(config, condition_id, template_arg, args):
+async def for_condition_to_code(
+    config: ConfigType,
+    condition_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     condition = await build_condition(
         config[CONF_CONDITION], cg.TemplateArguments(), []
     )
@@ -231,7 +274,12 @@ async def for_condition_to_code(config, condition_id, template_arg, args):
 @register_action(
     "delay", DelayAction, cv.templatable(cv.positive_time_period_milliseconds)
 )
-async def delay_action_to_code(config, action_id, template_arg, args):
+async def delay_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     var = cg.new_Pvariable(action_id, template_arg)
     await cg.register_component(var, {})
     template_ = await cg.templatable(config, args, cg.uint32)
@@ -256,10 +304,15 @@ async def delay_action_to_code(config, action_id, template_arg, args):
         cv.has_at_least_one_key(CONF_CONDITION, CONF_ANY, CONF_ALL),
     ),
 )
-async def if_action_to_code(config, action_id, template_arg, args):
+async def if_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     cond_conf = next(el for el in config if el in (CONF_ANY, CONF_ALL, CONF_CONDITION))
-    conditions = await build_condition(config[cond_conf], template_arg, args)
-    var = cg.new_Pvariable(action_id, template_arg, conditions)
+    condition = await build_condition(config[cond_conf], template_arg, args)
+    var = cg.new_Pvariable(action_id, template_arg, condition)
     if CONF_THEN in config:
         actions = await build_action_list(config[CONF_THEN], template_arg, args)
         cg.add(var.add_then(actions))
@@ -279,9 +332,14 @@ async def if_action_to_code(config, action_id, template_arg, args):
         }
     ),
 )
-async def while_action_to_code(config, action_id, template_arg, args):
-    conditions = await build_condition(config[CONF_CONDITION], template_arg, args)
-    var = cg.new_Pvariable(action_id, template_arg, conditions)
+async def while_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
+    condition = await build_condition(config[CONF_CONDITION], template_arg, args)
+    var = cg.new_Pvariable(action_id, template_arg, condition)
     actions = await build_action_list(config[CONF_THEN], template_arg, args)
     cg.add(var.add_then(actions))
     return var
@@ -297,7 +355,12 @@ async def while_action_to_code(config, action_id, template_arg, args):
         }
     ),
 )
-async def repeat_action_to_code(config, action_id, template_arg, args):
+async def repeat_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     var = cg.new_Pvariable(action_id, template_arg)
     count_template = await cg.templatable(config[CONF_COUNT], args, cg.uint32)
     cg.add(var.set_count(count_template))
@@ -320,9 +383,14 @@ _validate_wait_until = cv.maybe_simple_value(
 
 
 @register_action("wait_until", WaitUntilAction, _validate_wait_until)
-async def wait_until_action_to_code(config, action_id, template_arg, args):
-    conditions = await build_condition(config[CONF_CONDITION], template_arg, args)
-    var = cg.new_Pvariable(action_id, template_arg, conditions)
+async def wait_until_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
+    condition = await build_condition(config[CONF_CONDITION], template_arg, args)
+    var = cg.new_Pvariable(action_id, template_arg, condition)
     if CONF_TIMEOUT in config:
         template_ = await cg.templatable(config[CONF_TIMEOUT], args, cg.uint32)
         cg.add(var.set_timeout_value(template_))
@@ -331,7 +399,12 @@ async def wait_until_action_to_code(config, action_id, template_arg, args):
 
 
 @register_action("lambda", LambdaAction, cv.lambda_)
-async def lambda_action_to_code(config, action_id, template_arg, args):
+async def lambda_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     lambda_ = await cg.process_lambda(config, args, return_type=cg.void)
     return cg.new_Pvariable(action_id, template_arg, lambda_)
 
@@ -345,7 +418,12 @@ async def lambda_action_to_code(config, action_id, template_arg, args):
         }
     ),
 )
-async def component_update_action_to_code(config, action_id, template_arg, args):
+async def component_update_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     comp = await cg.get_variable(config[CONF_ID])
     return cg.new_Pvariable(action_id, template_arg, comp)
 
@@ -359,7 +437,12 @@ async def component_update_action_to_code(config, action_id, template_arg, args)
         }
     ),
 )
-async def component_suspend_action_to_code(config, action_id, template_arg, args):
+async def component_suspend_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     comp = await cg.get_variable(config[CONF_ID])
     return cg.new_Pvariable(action_id, template_arg, comp)
 
@@ -376,7 +459,12 @@ async def component_suspend_action_to_code(config, action_id, template_arg, args
         }
     ),
 )
-async def component_resume_action_to_code(config, action_id, template_arg, args):
+async def component_resume_action_to_code(
+    config: ConfigType,
+    action_id: ID,
+    template_arg: cg.TemplateArguments,
+    args: TemplateArgsType,
+) -> MockObj:
     comp = await cg.get_variable(config[CONF_ID])
     var = cg.new_Pvariable(action_id, template_arg, comp)
     if CONF_UPDATE_INTERVAL in config:
@@ -385,7 +473,9 @@ async def component_resume_action_to_code(config, action_id, template_arg, args)
     return var
 
 
-async def build_action(full_config, template_arg, args):
+async def build_action(
+    full_config: ConfigType, template_arg: cg.TemplateArguments, args: TemplateArgsType
+) -> MockObj:
     registry_entry, config = cg.extract_registry_entry_config(
         ACTION_REGISTRY, full_config
     )
@@ -394,15 +484,19 @@ async def build_action(full_config, template_arg, args):
     return await builder(config, action_id, template_arg, args)
 
 
-async def build_action_list(config, templ, arg_type):
-    actions = []
+async def build_action_list(
+    config: list[ConfigType], templ: cg.TemplateArguments, arg_type: TemplateArgsType
+) -> list[MockObj]:
+    actions: list[MockObj] = []
     for conf in config:
         action = await build_action(conf, templ, arg_type)
         actions.append(action)
     return actions
 
 
-async def build_condition(full_config, template_arg, args):
+async def build_condition(
+    full_config: ConfigType, template_arg: cg.TemplateArguments, args: TemplateArgsType
+) -> MockObj:
     registry_entry, config = cg.extract_registry_entry_config(
         CONDITION_REGISTRY, full_config
     )
@@ -411,15 +505,19 @@ async def build_condition(full_config, template_arg, args):
     return await builder(config, action_id, template_arg, args)
 
 
-async def build_condition_list(config, templ, args):
-    conditions = []
+async def build_condition_list(
+    config: ConfigType, templ: cg.TemplateArguments, args: TemplateArgsType
+) -> list[MockObj]:
+    conditions: list[MockObj] = []
     for conf in config:
         condition = await build_condition(conf, templ, args)
         conditions.append(condition)
     return conditions
 
 
-async def build_automation(trigger, args, config):
+async def build_automation(
+    trigger: MockObj, args: TemplateArgsType, config: ConfigType
+) -> MockObj:
     arg_types = [arg[0] for arg in args]
     templ = cg.TemplateArguments(*arg_types)
     obj = cg.new_Pvariable(config[CONF_AUTOMATION_ID], templ, trigger)

esphome/cpp_generator.py

@@ -1,5 +1,5 @@
 import abc
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 import inspect
 import math
 import re
@@ -13,7 +13,6 @@ from esphome.core import (
     HexInt,
     Lambda,
     Library,
-    TimePeriod,
     TimePeriodMicroseconds,
     TimePeriodMilliseconds,
     TimePeriodMinutes,
@@ -21,35 +20,11 @@ from esphome.core import (
     TimePeriodSeconds,
 )
 from esphome.helpers import cpp_string_escape, indent_all_but_first_and_last
+from esphome.types import Expression, SafeExpType, TemplateArgsType
 from esphome.util import OrderedDict
 from esphome.yaml_util import ESPHomeDataBase
 
 
-class Expression(abc.ABC):
-    __slots__ = ()
-
-    @abc.abstractmethod
-    def __str__(self):
-        """
-        Convert expression into C++ code
-        """
-
-
-SafeExpType = (
-    Expression
-    | bool
-    | str
-    | str
-    | int
-    | float
-    | TimePeriod
-    | type[bool]
-    | type[int]
-    | type[float]
-    | Sequence[Any]
-)
-
-
 class RawExpression(Expression):
     __slots__ = ("text",)
 
@@ -575,7 +550,7 @@ def Pvariable(id_: ID, rhs: SafeExpType, type_: "MockObj" = None) -> "MockObj":
     return obj
 
 
-def new_Pvariable(id_: ID, *args: SafeExpType) -> Pvariable:
+def new_Pvariable(id_: ID, *args: SafeExpType) -> "MockObj":
     """Declare a new pointer variable in the code generation by calling it's constructor
     with the given arguments.
 
@@ -681,7 +656,7 @@ async def get_variable_with_full_id(id_: ID) -> tuple[ID, "MockObj"]:
 
 async def process_lambda(
     value: Lambda,
-    parameters: list[tuple[SafeExpType, str]],
+    parameters: TemplateArgsType,
     capture: str = "=",
     return_type: SafeExpType = None,
 ) -> LambdaExpression | None:

esphome/types.py

@@ -1,8 +1,10 @@
 """This helper module tracks commonly used types in the esphome python codebase."""
 
-from typing import TypedDict
+import abc
+from collections.abc import Sequence
+from typing import Any, TypedDict
 
-from esphome.core import ID, EsphomeCore, Lambda
+from esphome.core import ID, EsphomeCore, Lambda, TimePeriod
 
 ConfigFragmentType = (
     str
@@ -20,6 +22,32 @@ CoreType = EsphomeCore
 ConfigPathType = str | int
 
 
+class Expression(abc.ABC):
+    __slots__ = ()
+
+    @abc.abstractmethod
+    def __str__(self):
+        """
+        Convert expression into C++ code
+        """
+
+
+SafeExpType = (
+    Expression
+    | bool
+    | str
+    | int
+    | float
+    | TimePeriod
+    | type[bool]
+    | type[int]
+    | type[float]
+    | Sequence[Any]
+)
+
+TemplateArgsType = list[tuple[SafeExpType, str]]
+
+
 class EntityMetadata(TypedDict):
     """Metadata stored for each entity to help with duplicate detection."""
 

esphome/util.py

@@ -1,19 +1,30 @@
 import collections
+from collections.abc import Callable
 import io
 import logging
 from pathlib import Path
 import re
 import subprocess
 import sys
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from esphome import const
 
 _LOGGER = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from esphome.config_validation import Schema
+    from esphome.cpp_generator import MockObjClass
+
 
 class RegistryEntry:
-    def __init__(self, name, fun, type_id, schema):
+    def __init__(
+        self,
+        name: str,
+        fun: Callable[..., Any],
+        type_id: "MockObjClass",
+        schema: "Schema",
+    ):
         self.name = name
         self.fun = fun
         self.type_id = type_id
@@ -38,8 +49,8 @@ class Registry(dict[str, RegistryEntry]):
         self.base_schema = base_schema or {}
         self.type_id_key = type_id_key
 
-    def register(self, name, type_id, schema):
-        def decorator(fun):
+    def register(self, name: str, type_id: "MockObjClass", schema: "Schema"):
+        def decorator(fun: Callable[..., Any]):
             self[name] = RegistryEntry(name, fun, type_id, schema)
             return fun
 
@@ -47,8 +58,8 @@ class Registry(dict[str, RegistryEntry]):
 
 
 class SimpleRegistry(dict):
-    def register(self, name, data):
-        def decorator(fun):
+    def register(self, name: str, data: Any):
+        def decorator(fun: Callable[..., Any]):
             self[name] = (fun, data)
             return fun