[core] Reduce memory footprint of esphome upload (#17684)

esphome/__main__.py

@@ -16,14 +16,10 @@ import sys
 import time
 from typing import Protocol
 
-import argcomplete
-
 # Note: Do not import modules from esphome.components here, as this would
 # cause them to be loaded before external components are processed, resulting
 # in the built-in version being used instead of the external component one.
 from esphome import const
-import esphome.codegen as cg
-from esphome.config import iter_component_configs, read_config, strip_default_ids
 from esphome.const import (
     ALLOWED_NAME_CHARS,
     ARGUMENT_HELP_DEVICE,
@@ -704,6 +700,8 @@ def run_miniterm(config: ConfigType, port: str, args) -> int:
 
 
 def _wrap_to_code(name, comp, yaml_util):
+    import esphome.codegen as cg
+
     coro = coroutine(comp.to_code)
 
     @functools.wraps(comp.to_code)
@@ -739,6 +737,7 @@ def write_cpp(config: ConfigType) -> int:
 
 def generate_cpp_contents(config: ConfigType) -> None:
     from esphome import yaml_util
+    from esphome.config import iter_component_configs
 
     _LOGGER.info("Generating C++ source...")
 
@@ -1464,6 +1463,7 @@ def command_wizard(args: ArgsProtocol) -> int | None:
 
 def command_config(args: ArgsProtocol, config: ConfigType) -> int | None:
     from esphome import yaml_util
+    from esphome.config import strip_default_ids
 
     if getattr(args, "no_defaults", False):
         user_config = getattr(config, "user_config", None)
@@ -2498,7 +2498,12 @@ def parse_args(argv):
     # a deprecation warning).
     arguments = argv[1:]
 
-    argcomplete.autocomplete(parser)
+    # argcomplete only does anything when the shell-completion machinery
+    # invokes us with _ARGCOMPLETE set; skip the import otherwise.
+    if "_ARGCOMPLETE" in os.environ:
+        import argcomplete
+
+        argcomplete.autocomplete(parser)
 
     if len(arguments) > 0 and arguments[0] in SIMPLE_CONFIG_ACTIONS:
         args, unknown_args = parser.parse_known_args(arguments)
@@ -2597,6 +2602,8 @@ def run_esphome(argv):
             )
 
     if config is None:
+        from esphome.config import read_config
+
         config = read_config(
             command_line_substitutions,
             skip_external_update=skip_external,

script/import_time_budget.json

@@ -1,5 +1,5 @@
 {
   "target_module": "esphome.__main__",
   "margin_pct": 20,
-  "cumulative_us": 95000
+  "cumulative_us": 37200
 }

tests/unit_tests/test_compiled_config.py

@@ -220,7 +220,7 @@ def test_run_esphome_upload_and_logs_use_cache_when_fresh(
 
     with (
         caplog.at_level("INFO", logger="esphome.__main__"),
-        patch("esphome.__main__.read_config") as mock_read,
+        patch("esphome.config.read_config") as mock_read,
         patch.dict("esphome.__main__.POST_CONFIG_ACTIONS", {command: _stub}),
     ):
         assert run_esphome(["esphome", command, str(fresh_cache_files)]) == 0
@@ -242,7 +242,7 @@ def test_run_esphome_upload_and_logs_fall_back_when_no_cache(
     yaml_path.write_text("esphome:\n  name: lite_test\n")
 
     with (
-        patch("esphome.__main__.read_config", return_value=None) as mock_read,
+        patch("esphome.config.read_config", return_value=None) as mock_read,
         patch.dict(
             "esphome.__main__.POST_CONFIG_ACTIONS",
             {command: lambda args, config: 0},
@@ -266,7 +266,7 @@ def test_run_esphome_upload_does_not_refresh_cache_without_sidecar(
 
     with (
         patch(
-            "esphome.__main__.read_config",
+            "esphome.config.read_config",
             return_value={"esphome": {"name": "lite_test"}},
         ),
         patch("esphome.compiled_config.save_compiled_config") as mock_save,
@@ -299,7 +299,7 @@ def test_run_esphome_upload_and_logs_refresh_cache_on_fallback(
     fresh_config = {"esphome": {"name": "lite_test"}, "logger": {}}
 
     with (
-        patch("esphome.__main__.read_config", return_value=fresh_config),
+        patch("esphome.config.read_config", return_value=fresh_config),
         patch(
             "esphome.compiled_config.save_compiled_config", wraps=save_compiled_config
         ) as mock_save,
@@ -322,7 +322,7 @@ def test_run_esphome_upload_with_substitution_does_not_refresh_cache(
     """`-s` substitutions skip the cache on both read and write -- saving
     here would clobber the cache with a substitution-specific config."""
     with (
-        patch("esphome.__main__.read_config", return_value={"esphome": {}}),
+        patch("esphome.config.read_config", return_value={"esphome": {}}),
         patch("esphome.compiled_config.save_compiled_config") as mock_save,
         patch.dict(
             "esphome.__main__.POST_CONFIG_ACTIONS",
@@ -341,7 +341,7 @@ def test_run_esphome_compile_does_not_refresh_cache_via_fallback(
     upload/logs fallback path -- the fallback save would skip the
     storage_should_clean check."""
     with (
-        patch("esphome.__main__.read_config", return_value={"esphome": {}}),
+        patch("esphome.config.read_config", return_value={"esphome": {}}),
         patch("esphome.compiled_config.save_compiled_config") as mock_save,
         patch.dict(
             "esphome.__main__.POST_CONFIG_ACTIONS",
@@ -360,7 +360,7 @@ def test_run_esphome_upload_with_substitution_skips_cache(
     against the prior substitution set, so reusing it would silently
     ignore the override."""
     with (
-        patch("esphome.__main__.read_config", return_value=None) as mock_read,
+        patch("esphome.config.read_config", return_value=None) as mock_read,
         patch.dict(
             "esphome.__main__.POST_CONFIG_ACTIONS",
             {"upload": lambda args, config: 0},
@@ -374,7 +374,7 @@ def test_run_esphome_upload_with_substitution_skips_cache(
 def test_run_esphome_compile_does_not_use_cache(fresh_cache_files: Path) -> None:
     """The compile subcommand always re-validates -- it's what writes the cache."""
     with (
-        patch("esphome.__main__.read_config", return_value=None) as mock_read,
+        patch("esphome.config.read_config", return_value=None) as mock_read,
         patch.dict(
             "esphome.__main__.POST_CONFIG_ACTIONS",
             {"compile": lambda args, config: 0},

tests/unit_tests/test_lazy_imports.py

@@ -0,0 +1,53 @@
+"""Guard the lazy-import contract of ``esphome.__main__``.
+
+Every ``esphome`` invocation pays for whatever ``esphome.__main__``
+imports at module level before the requested command runs. The
+dashboard and device-builder spawn one ``esphome upload`` subprocess
+per device, so keeping validation/codegen machinery out of the
+top-level import directly lowers the RAM cost of each concurrent
+upload (the upload/logs fast path in ``esphome.compiled_config``
+never needs them).
+
+``script/check_import_time.py`` budgets import *time* in CI; this
+test pins down *which* heavy modules must stay out entirely.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+# Modules that must only load for the commands that actually use them
+# (compile/config validation, shell completion), never from a bare
+# ``import esphome.__main__``.
+HEAVY_MODULES = (
+    "argcomplete",
+    "esphome.codegen",
+    "esphome.config",
+    "esphome.config_validation",
+    "esphome.cpp_generator",
+    "esphome.loader",
+    "voluptuous",
+)
+
+
+def test_main_module_does_not_import_heavy_modules() -> None:
+    """A bare ``import esphome.__main__`` must not drag in validation/codegen."""
+    check = (
+        "import sys; import esphome.__main__; "
+        f"leaked = [m for m in {HEAVY_MODULES!r} if m in sys.modules]; "
+        "print(','.join(leaked))"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", check],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    leaked = result.stdout.strip()
+    assert not leaked, (
+        f"esphome.__main__ imports heavy modules at top level: {leaked}. "
+        "Import them lazily inside the command that needs them instead; "
+        "every esphome invocation (including each parallel dashboard "
+        "upload subprocess) pays for top-level imports."
+    )

tests/unit_tests/test_main.py

@@ -618,7 +618,7 @@ def test_command_config__no_defaults_skips_strip_default_ids(
     validated.user_config = {"sensor": [{"name": "x"}]}
 
     with patch(
-        "esphome.__main__.strip_default_ids", side_effect=AssertionError
+        "esphome.config.strip_default_ids", side_effect=AssertionError
     ) as mock_strip:
         result = command_config(args, validated)
 
@@ -6201,7 +6201,7 @@ def test_run_esphome_bundle_detection(tmp_path: Path) -> None:
             "esphome.bundle.prepare_bundle_for_compile",
             return_value=extracted_yaml,
         ) as mock_prepare,
-        patch("esphome.__main__.read_config", return_value=None),
+        patch("esphome.config.read_config", return_value=None),
     ):
         result = run_esphome(["esphome", "compile", str(bundle_path)])
 
@@ -6219,7 +6219,7 @@ def test_run_esphome_non_bundle_skips_extraction(tmp_path: Path) -> None:
     with (
         patch("esphome.bundle.is_bundle_path", return_value=False) as mock_is_bundle,
         patch("esphome.bundle.prepare_bundle_for_compile") as mock_prepare,
-        patch("esphome.__main__.read_config", return_value=None),
+        patch("esphome.config.read_config", return_value=None),
     ):
         result = run_esphome(["esphome", "compile", str(yaml_file)])
 
@@ -6247,7 +6247,7 @@ def test_run_esphome_skip_external_update_per_command(
     yaml_file = tmp_path / "device.yaml"
     yaml_file.write_text("esphome:\n  name: test\n")
 
-    with patch("esphome.__main__.read_config", return_value=None) as mock_read:
+    with patch("esphome.config.read_config", return_value=None) as mock_read:
         run_esphome(["esphome", command, str(yaml_file)])
 
     mock_read.assert_called_once()
@@ -6405,6 +6405,23 @@ def test_parse_args_logs_states() -> None:
     assert args.states is True
 
 
+def test_parse_args_argcomplete_only_runs_when_completing() -> None:
+    """Only import and invoke argcomplete when _ARGCOMPLETE is set.
+
+    The shell-completion machinery sets _ARGCOMPLETE when it invokes the
+    CLI; a normal invocation must skip the import entirely so every
+    esphome subprocess (e.g. parallel dashboard uploads) avoids paying
+    for it.
+    """
+    fake_argcomplete = MagicMock()
+    with (
+        patch.dict(os.environ, {"_ARGCOMPLETE": "1"}),
+        patch.dict(sys.modules, {"argcomplete": fake_argcomplete}),
+    ):
+        parse_args(["esphome", "version"])
+    fake_argcomplete.autocomplete.assert_called_once()
+
+
 def test_should_subscribe_states_default() -> None:
     """Test that states are shown by default when nothing is set."""
     from esphome.__main__ import _should_subscribe_states