Skip to content

Commit

Permalink
Speed up the frame helper
Browse files Browse the repository at this point in the history 
This is same solution we have in system_log

Inspecting the stack got a lot more expensive in python 3.11
see python/cpython#92041

see getsentry/sentry#63965
  • Loading branch information
@bdraco
bdraco committed Mar 7, 2024
1 parent 67a1776 commit ca80f5d
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 71 deletions.
32 changes: 21 additions & 11 deletions homeassistant/helpers/frame.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import functools
import logging
import sys
from traceback import FrameSummary, extract_stack
from types import FrameType
from typing import Any, TypeVar, cast

from homeassistant.core import HomeAssistant, async_get_hass
Expand All @@ -28,7 +28,7 @@ class IntegrationFrame:
"""Integration frame container."""

custom_integration: bool
frame: FrameSummary
frame: FrameType
integration: str
module: str | None
relative_filename: str
Expand All @@ -54,19 +54,27 @@ def get_integration_logger(fallback_name: str) -> logging.Logger:
return logging.getLogger(logger_name)


def get_current_frame() -> FrameType:
"""Return the current frame."""
return sys._getframe(1) # pylint: disable=protected-access


def get_integration_frame(exclude_integrations: set | None = None) -> IntegrationFrame:
"""Return the frame, integration and integration path of the current stack frame."""
found_frame = None
if not exclude_integrations:
exclude_integrations = set()

for frame in reversed(extract_stack()):
frame: FrameType | None = get_current_frame()
while frame is not None:
filename = frame.f_code.co_filename

for path in ("custom_components/", "homeassistant/components/"):
try:
index = frame.filename.index(path)
index = filename.index(path)
start = index + len(path)
end = frame.filename.index("/", start)
integration = frame.filename[start:end]
end = filename.index("/", start)
integration = filename[start:end]
if integration not in exclude_integrations:
found_frame = frame

Expand All @@ -77,14 +85,16 @@ def get_integration_frame(exclude_integrations: set | None = None) -> Integratio
if found_frame is not None:
break

frame = frame.f_back

if found_frame is None:
raise MissingIntegrationFrame

found_module: str | None = None
for module, module_obj in dict(sys.modules).items():
if not hasattr(module_obj, "__file__"):
continue
if module_obj.__file__ == found_frame.filename:
if module_obj.__file__ == found_frame.f_code.co_filename:
found_module = module
break

Expand All @@ -93,7 +103,7 @@ def get_integration_frame(exclude_integrations: set | None = None) -> Integratio
frame=found_frame,
integration=integration,
module=found_module,
relative_filename=found_frame.filename[index:],
relative_filename=found_frame.f_code.co_filename[index:],
)


Expand Down Expand Up @@ -139,7 +149,7 @@ def _report_integration(
"""
found_frame = integration_frame.frame
# Keep track of integrations already reported to prevent flooding
key = f"{found_frame.filename}:{found_frame.lineno}"
key = f"{found_frame.f_code.co_filename}:{found_frame.f_code.co_firstlineno}"
if key in _REPORTED_INTEGRATIONS:
return
_REPORTED_INTEGRATIONS.add(key)
Expand All @@ -160,8 +170,8 @@ def _report_integration(
integration_frame.integration,
what,
integration_frame.relative_filename,
found_frame.lineno,
(found_frame.line or "?").strip(),
found_frame.f_code.co_firstlineno,
(str(found_frame) or "?").strip(),
report_issue,
)

Expand Down
18 changes: 17 additions & 1 deletion tests/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
import pathlib
import threading
import time
from types import ModuleType
from types import FrameType, ModuleType
from typing import Any, NoReturn, TypeVar
from unittest.mock import AsyncMock, Mock, patch

Expand Down Expand Up @@ -1596,3 +1596,19 @@ def help_test_all(module: ModuleType) -> None:
assert set(module.__all__) == {
itm for itm in module.__dir__() if not itm.startswith("_")
}


def extract_stack_to_frame(extract_stack: list[Mock]) -> FrameType:
"""Convert an extract stack to a frame list."""
stack = list(extract_stack)
for frame in stack:
frame.f_back = None
frame.f_code.co_filename = frame.filename

top_frame = stack.pop()
current_frame = top_frame
while stack and (next_frame := stack.pop()):
current_frame.f_back = next_frame
current_frame = next_frame

return top_frame
31 changes: 17 additions & 14 deletions tests/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@
init_recorder_component,
mock_storage,
patch_yaml_files,
extract_stack_to_frame,
)
from .test_util.aiohttp import ( # noqa: E402, isort:skip
AiohttpClientMocker,
Expand Down Expand Up @@ -1588,20 +1589,22 @@ def mock_integration_frame() -> Generator[Mock, None, None]:
line="self.light.is_on",
)
with patch(
"homeassistant.helpers.frame.extract_stack",
return_value=[
Mock(
filename="/home/paulus/homeassistant/core.py",
lineno="23",
line="do_something()",
),
correct_frame,
Mock(
filename="/home/paulus/aiohue/lights.py",
lineno="2",
line="something()",
),
],
"homeassistant.helpers.frame.get_current_frame",
return_value=extract_stack_to_frame(
[
Mock(
filename="/home/paulus/homeassistant/core.py",
lineno="23",
line="do_something()",
),
correct_frame,
Mock(
filename="/home/paulus/aiohue/lights.py",
lineno="2",
line="something()",
),
]
),
):
yield correct_frame

Expand Down
98 changes: 53 additions & 45 deletions tests/helpers/test_frame.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
from homeassistant.core import HomeAssistant
from homeassistant.helpers import frame

from tests.common import extract_stack_to_frame


async def test_extract_frame_integration(
caplog: pytest.LogCaptureFixture, mock_integration_frame: Mock
Expand Down Expand Up @@ -68,25 +70,27 @@ async def test_extract_frame_integration_with_excluded_integration(
line="self.light.is_on",
)
with patch(
"homeassistant.helpers.frame.extract_stack",
return_value=[
Mock(
filename="/home/dev/homeassistant/core.py",
lineno="23",
line="do_something()",
),
correct_frame,
Mock(
filename="/home/dev/homeassistant/components/zeroconf/usage.py",
lineno="23",
line="self.light.is_on",
),
Mock(
filename="/home/dev/mdns/lights.py",
lineno="2",
line="something()",
),
],
"homeassistant.helpers.frame.get_current_frame",
return_value=extract_stack_to_frame(
[
Mock(
filename="/home/dev/homeassistant/core.py",
lineno="23",
line="do_something()",
),
correct_frame,
Mock(
filename="/home/dev/homeassistant/components/zeroconf/usage.py",
lineno="23",
line="self.light.is_on",
),
Mock(
filename="/home/dev/mdns/lights.py",
lineno="2",
line="something()",
),
]
),
):
integration_frame = frame.get_integration_frame(
exclude_integrations={"zeroconf"}
Expand All @@ -104,19 +108,21 @@ async def test_extract_frame_integration_with_excluded_integration(
async def test_extract_frame_no_integration(caplog: pytest.LogCaptureFixture) -> None:
"""Test extracting the current frame without integration context."""
with patch(
"homeassistant.helpers.frame.extract_stack",
return_value=[
Mock(
filename="/home/paulus/homeassistant/core.py",
lineno="23",
line="do_something()",
),
Mock(
filename="/home/paulus/aiohue/lights.py",
lineno="2",
line="something()",
),
],
"homeassistant.helpers.frame.get_current_frame",
return_value=extract_stack_to_frame(
[
Mock(
filename="/home/paulus/homeassistant/core.py",
lineno="23",
line="do_something()",
),
Mock(
filename="/home/paulus/aiohue/lights.py",
lineno="2",
line="something()",
),
]
),
), pytest.raises(frame.MissingIntegrationFrame):
frame.get_integration_frame()

Expand All @@ -126,19 +132,21 @@ async def test_get_integration_logger_no_integration(
) -> None:
"""Test getting fallback logger without integration context."""
with patch(
"homeassistant.helpers.frame.extract_stack",
return_value=[
Mock(
filename="/home/paulus/homeassistant/core.py",
lineno="23",
line="do_something()",
),
Mock(
filename="/home/paulus/aiohue/lights.py",
lineno="2",
line="something()",
),
],
"homeassistant.helpers.frame.get_current_frame",
return_value=extract_stack_to_frame(
[
Mock(
filename="/home/paulus/homeassistant/core.py",
lineno="23",
line="do_something()",
),
Mock(
filename="/home/paulus/aiohue/lights.py",
lineno="2",
line="something()",
),
]
),
):
logger = frame.get_integration_logger(__name__)

Expand Down

0 comments on commit ca80f5d

Please sign in to comment.