Fix typing errors involving handle_internal_errors (#885)

Author: alexmojaki

logfire/_internal/config.py

@@ -317,22 +317,22 @@ def configure(  # noqa: D417
     """
     from .. import DEFAULT_LOGFIRE_INSTANCE, Logfire
 
-    processors = deprecated_kwargs.pop('processors', None)  # type: ignore
+    processors = deprecated_kwargs.pop('processors', None)
     if processors is not None:  # pragma: no cover
         raise ValueError(
             'The `processors` argument has been replaced by `additional_span_processors`. '
             'Set `send_to_logfire=False` to disable the default processor.'
         )
 
-    metric_readers = deprecated_kwargs.pop('metric_readers', None)  # type: ignore
+    metric_readers = deprecated_kwargs.pop('metric_readers', None)
     if metric_readers is not None:  # pragma: no cover
         raise ValueError(
             'The `metric_readers` argument has been replaced by '
             '`metrics=logfire.MetricsOptions(additional_readers=[...])`. '
             'Set `send_to_logfire=False` to disable the default metric reader.'
         )
 
-    collect_system_metrics = deprecated_kwargs.pop('collect_system_metrics', None)  # type: ignore
+    collect_system_metrics = deprecated_kwargs.pop('collect_system_metrics', None)
     if collect_system_metrics is False:
         raise ValueError(
             'The `collect_system_metrics` argument has been removed. System metrics are no longer collected by default.'
@@ -343,8 +343,8 @@ def configure(  # noqa: D417
             'The `collect_system_metrics` argument has been removed. Use `logfire.instrument_system_metrics()` instead.'
         )
 
-    scrubbing_callback = deprecated_kwargs.pop('scrubbing_callback', None)  # type: ignore
-    scrubbing_patterns = deprecated_kwargs.pop('scrubbing_patterns', None)  # type: ignore
+    scrubbing_callback = deprecated_kwargs.pop('scrubbing_callback', None)
+    scrubbing_patterns = deprecated_kwargs.pop('scrubbing_patterns', None)
     if scrubbing_callback or scrubbing_patterns:
         if scrubbing is not None:
             raise ValueError(
@@ -357,7 +357,7 @@ def configure(  # noqa: D417
         )
         scrubbing = ScrubbingOptions(callback=scrubbing_callback, extra_patterns=scrubbing_patterns)  # type: ignore
 
-    project_name = deprecated_kwargs.pop('project_name', None)  # type: ignore
+    project_name = deprecated_kwargs.pop('project_name', None)
     if project_name is not None:
         warnings.warn(
             'The `project_name` argument is deprecated and not needed.',
@@ -377,15 +377,15 @@ def configure(  # noqa: D417
                 'Use `sampling=logfire.SamplingOptions(head=...)` instead.',
             )
 
-    show_summary = deprecated_kwargs.pop('show_summary', None)  # type: ignore
+    show_summary = deprecated_kwargs.pop('show_summary', None)
     if show_summary is not None:  # pragma: no cover
         warnings.warn(
             'The `show_summary` argument is deprecated. '
             'Use `console=False` or `console=logfire.ConsoleOptions(show_project_link=False)` instead.',
         )
 
     for key in ('base_url', 'id_generator', 'ns_timestamp_generator'):
-        value: Any = deprecated_kwargs.pop(key, None)  # type: ignore
+        value: Any = deprecated_kwargs.pop(key, None)
         if value is None:
             continue
         if advanced is not None:
@@ -397,7 +397,7 @@ def configure(  # noqa: D417
             stacklevel=2,
         )
 
-    additional_metric_readers: Any = deprecated_kwargs.pop('additional_metric_readers', None)  # type: ignore
+    additional_metric_readers: Any = deprecated_kwargs.pop('additional_metric_readers', None)
     if additional_metric_readers:
         if metrics is not None:
             raise ValueError(
@@ -410,7 +410,7 @@ def configure(  # noqa: D417
         )
         metrics = MetricsOptions(additional_readers=additional_metric_readers)
 
-    pydantic_plugin: Any = deprecated_kwargs.pop('pydantic_plugin', None)  # type: ignore
+    pydantic_plugin: Any = deprecated_kwargs.pop('pydantic_plugin', None)
     if pydantic_plugin is not None:
         warnings.warn(
             'The `pydantic_plugin` argument is deprecated. Use `logfire.instrument_pydantic()` instead.',
@@ -994,7 +994,7 @@ def check_token():
             if hasattr(os, 'register_at_fork'):  # pragma: no branch
 
                 def fix_pid():  # pragma: no cover
-                    with handle_internal_errors():
+                    with handle_internal_errors:
                         new_resource = resource.merge(Resource({ResourceAttributes.PROCESS_PID: os.getpid()}))
                         tracer_provider._resource = new_resource  # type: ignore
                         meter_provider._resource = new_resource  # type: ignore

logfire/_internal/integrations/fastapi.py

@@ -167,7 +167,7 @@ async def solve_dependencies(self, request: Request | WebSocket, original: Await
             return await original
 
         with self.logfire_instance.span('FastAPI arguments') as span:
-            with handle_internal_errors():
+            with handle_internal_errors:
                 if isinstance(request, Request):  # pragma: no branch
                     span.set_attribute(SpanAttributes.HTTP_METHOD, request.method)
                 route: APIRoute | APIWebSocketRoute | None = request.scope.get('route')
@@ -181,7 +181,7 @@ async def solve_dependencies(self, request: Request | WebSocket, original: Await
 
             result: Any = await original
 
-            with handle_internal_errors():
+            with handle_internal_errors:
                 solved_values: dict[str, Any]
                 solved_errors: list[Any]
 

logfire/_internal/integrations/httpx.py

@@ -322,7 +322,7 @@ def make_request_hook(hook: RequestHook | None, capture_headers: bool, capture_b
         return None
 
     def new_hook(span: Span, request: RequestInfo) -> None:
-        with handle_internal_errors():
+        with handle_internal_errors:
             request = capture_request(span, request, capture_headers, capture_body)
             run_hook(hook, span, request)
 
@@ -338,7 +338,7 @@ def make_async_request_hook(
         return None
 
     async def new_hook(span: Span, request: RequestInfo) -> None:
-        with handle_internal_errors():
+        with handle_internal_errors:
             request = capture_request(span, request, should_capture_headers, should_capture_body)
             await run_async_hook(hook, span, request)
 
@@ -355,7 +355,7 @@ def make_response_hook(
         return None
 
     def new_hook(span: Span, request: RequestInfo, response: ResponseInfo) -> None:
-        with handle_internal_errors():
+        with handle_internal_errors:
             request, response = capture_response(
                 span,
                 request,
@@ -380,7 +380,7 @@ def make_async_response_hook(
         return None
 
     async def new_hook(span: Span, request: RequestInfo, response: ResponseInfo) -> None:
-        with handle_internal_errors():
+        with handle_internal_errors:
             request, response = capture_response(
                 span,
                 request,

logfire/_internal/integrations/psycopg.py

@@ -41,6 +41,7 @@ def instrument_psycopg(
                 instrument_psycopg(logfire_instance, package, **kwargs)
         return
     elif conn_or_module in PACKAGE_NAMES:
+        assert isinstance(conn_or_module, str)
         _instrument_psycopg(logfire_instance, name=conn_or_module, **kwargs)
         return
     elif isinstance(conn_or_module, ModuleType):

logfire/_internal/main.py

@@ -662,7 +662,7 @@ def log(
                 Set to `True` to use the currently handled exception.
             console_log: Whether to log to the console, defaults to `True`.
         """
-        with handle_internal_errors():
+        with handle_internal_errors:
             stack_info = get_user_stack_info()
 
             attributes = attributes or {}
@@ -2107,7 +2107,7 @@ def __init__(self, span: trace_api.Span) -> None:
     def __enter__(self) -> FastLogfireSpan:
         return self
 
-    @handle_internal_errors()
+    @handle_internal_errors
     def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None:
         context_api.detach(self._token)
         self._span.__exit__(exc_type, exc_value, traceback)
@@ -2139,7 +2139,7 @@ def __getattr__(self, name: str) -> Any:
             return getattr(self._span, name)
 
     def __enter__(self) -> LogfireSpan:
-        with handle_internal_errors():
+        with handle_internal_errors:
             if self._span is None:  # pragma: no branch
                 self._span = self._tracer.start_span(
                     name=self._span_name,
@@ -2152,7 +2152,7 @@ def __enter__(self) -> LogfireSpan:
 
         return self
 
-    @handle_internal_errors()
+    @handle_internal_errors
     def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any) -> None:
         if self._token is None:  # pragma: no cover
             return
@@ -2161,7 +2161,7 @@ def __exit__(self, exc_type: type[BaseException] | None, exc_value: BaseExceptio
         context_api.detach(self._token)
         self._token = None
         if self._span.is_recording():
-            with handle_internal_errors():
+            with handle_internal_errors:
                 if self._added_attributes:
                     self._span.set_attribute(
                         ATTRIBUTES_JSON_SCHEMA_KEY, attributes_json_schema(self._json_schema_properties)
@@ -2177,7 +2177,7 @@ def tags(self) -> tuple[str, ...]:
         return self._get_attribute(ATTRIBUTES_TAGS_KEY, ())
 
     @tags.setter
-    @handle_internal_errors()
+    @handle_internal_errors
     def tags(self, new_tags: Sequence[str]) -> None:
         """Set or add tags to the span."""
         if isinstance(new_tags, str):
@@ -2192,7 +2192,7 @@ def message(self) -> str:
     def message(self, message: str):
         self._set_attribute(ATTRIBUTES_MESSAGE_KEY, message)
 
-    @handle_internal_errors()
+    @handle_internal_errors
     def set_attribute(self, key: str, value: Any) -> None:
         """Sets an attribute on the span.
 
@@ -2247,7 +2247,7 @@ def record_exception(
     def is_recording(self) -> bool:
         return self._span is not None and self._span.is_recording()
 
-    @handle_internal_errors()
+    @handle_internal_errors
     def set_level(self, level: LevelName | int):
         """Set the log level of this span."""
         attributes = log_level_attributes(level)

logfire/_internal/tracer.py

@@ -345,7 +345,7 @@ def get_sample_rate_from_attributes(attributes: otel_types.Attributes) -> float
     return cast('float | None', attributes.get(ATTRIBUTES_SAMPLE_RATE_KEY))
 
 
-@handle_internal_errors()
+@handle_internal_errors
 def record_exception(
     span: trace_api.Span,
     exception: BaseException,

logfire/_internal/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import inspect
 import json
 import logging
@@ -12,7 +13,19 @@
 from pathlib import Path
 from time import time
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Sequence, Tuple, TypedDict, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Sequence,
+    Tuple,
+    TypedDict,
+    TypeVar,
+    Union,
+)
 
 from opentelemetry import context, trace as trace_api
 from opentelemetry.sdk.resources import Resource
@@ -27,13 +40,17 @@
 from logfire._internal.ulid import ulid
 
 if TYPE_CHECKING:
+    from typing import ParamSpec
+
     from packaging.version import Version
 
     SysExcInfo = Union[tuple[type[BaseException], BaseException, TracebackType | None], tuple[None, None, None]]
     """
     The return type of sys.exc_info(): exc_type, exc_val, exc_tb.
     """
 
+    P = ParamSpec('P')
+
 T = TypeVar('T')
 
 JsonValue = Union[int, float, str, bool, None, List['JsonValue'], Tuple['JsonValue', ...], 'JsonDict']
@@ -280,20 +297,9 @@ def _internal_error_exc_info() -> SysExcInfo:
     original_exc_info: SysExcInfo = sys.exc_info()
     exc_type, exc_val, original_tb = original_exc_info
     try:
-        # First remove redundant frames already in the traceback about where the error was raised.
         tb = original_tb
-        if tb and tb.tb_frame and tb.tb_frame.f_code is _HANDLE_INTERNAL_ERRORS_CODE:
-            # Skip the 'yield' line in _handle_internal_errors
-            tb = tb.tb_next
-
-        if (
-            tb
-            and tb.tb_frame
-            and tb.tb_frame.f_code.co_filename == contextmanager.__code__.co_filename
-            and tb.tb_frame.f_code.co_name == 'inner'
-        ):
-            # Skip the 'inner' function frame when handle_internal_errors is used as a decorator.
-            # It looks like `return func(*args, **kwds)`
+        if tb and tb.tb_frame and tb.tb_frame.f_code is _HANDLE_INTERNAL_ERRORS_WRAPPER_CODE:
+            # Skip the redundant `with self:` line in the traceback about where the error was raised.
             tb = tb.tb_next
 
         # Now add useful outer frames that give context, but skipping frames that are just about handling the error.
@@ -307,18 +313,16 @@ def _internal_error_exc_info() -> SysExcInfo:
             frame = frame.f_back
             assert frame
 
-            if frame.f_code is _HANDLE_INTERNAL_ERRORS_CODE:
-                # Skip the line in _handle_internal_errors that calls log_internal_error
-                frame = frame.f_back
-                # Skip the frame defining the _handle_internal_errors context manager
-                assert frame and frame.f_code.co_name == '__exit__'
+            if frame.f_code is _HANDLE_INTERNAL_ERRORS_EXIT_CODE:
+                # Skip the `log_internal_error()` call in `__exit__`.
                 frame = frame.f_back
                 assert frame
-                # Skip the frame calling the context manager, on the `with` line.
-                frame = frame.f_back
-            else:
-                # `log_internal_error()` was called directly, so just skip that frame. No context manager stuff.
-                frame = frame.f_back
+
+            # Now skip the line that is either:
+            # - A direct call to `log_internal_error`
+            # - `with self:` in HandleInternalErrors.__call__.wrapper
+            # - `with handle_internal_errors:`
+            frame = frame.f_back
 
         # Now add all remaining frames from internal logfire code.
         while frame and not is_user_code(frame.f_code):
@@ -340,15 +344,29 @@ def _internal_error_exc_info() -> SysExcInfo:
         return original_exc_info
 
 
-@contextmanager
-def handle_internal_errors():
-    try:
-        yield
-    except Exception:
-        log_internal_error()
+class HandleInternalErrors:
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> bool | None:
+        if isinstance(exc_val, Exception):
+            log_internal_error()
+            return True
+
+    def __call__(self, func: Callable[P, T]) -> Callable[P, T]:
+        @functools.wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            with self:
+                return func(*args, **kwargs)
+
+        return wrapper
+
 
+handle_internal_errors = HandleInternalErrors()
 
-_HANDLE_INTERNAL_ERRORS_CODE = inspect.unwrap(handle_internal_errors).__code__
+_HANDLE_INTERNAL_ERRORS_WRAPPER_CODE = handle_internal_errors(log_internal_error).__code__
+assert _HANDLE_INTERNAL_ERRORS_WRAPPER_CODE.co_name == 'wrapper'
+_HANDLE_INTERNAL_ERRORS_EXIT_CODE = HandleInternalErrors.__exit__.__code__
 
 
 def maybe_capture_server_headers(capture: bool):

tests/import_used_for_tests/internal_error_handling/internal_logfire_code_example.py

@@ -11,13 +11,13 @@ def inner2():
     inner1()
 
 
-@handle_internal_errors()
+@handle_internal_errors
 def using_decorator():
     inner2()
 
 
 def using_context_manager():
-    with handle_internal_errors():
+    with handle_internal_errors:
         inner2()