Add non-user-code-prefix stuff (#829)

Co-authored-by: Alex Hall <[email protected]>

Author: dmontagu

logfire-api/logfire_api/__init__.py

@@ -193,6 +193,8 @@ def loguru_handler() -> dict[str, Any]: ...
         def no_auto_trace(x):
             return x
 
+        def add_non_user_code_prefix(*args, **kwargs) -> None: ...
+
         @contextmanager
         def suppress_instrumentation():
             yield

logfire-api/logfire_api/__init__.pyi

@@ -4,14 +4,15 @@ from ._internal.config import AdvancedOptions as AdvancedOptions, CodeSource as
 from ._internal.constants import LevelName as LevelName
 from ._internal.main import Logfire as Logfire, LogfireSpan as LogfireSpan
 from ._internal.scrubbing import ScrubMatch as ScrubMatch, ScrubbingOptions as ScrubbingOptions
+from ._internal.stack_info import add_non_user_code_prefix as add_non_user_code_prefix
 from ._internal.utils import suppress_instrumentation as suppress_instrumentation
 from .integrations.logging import LogfireLoggingHandler as LogfireLoggingHandler
 from .integrations.structlog import LogfireProcessor as StructlogProcessor
 from .version import VERSION as VERSION
 from logfire.sampling import SamplingOptions as SamplingOptions
 from typing import Any
 
-__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'warning', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_sqlite3', 'instrument_aws_lambda', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'suppress_scopes', 'shutdown', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
+__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'warning', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_sqlite3', 'instrument_aws_lambda', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'suppress_scopes', 'shutdown', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'add_non_user_code_prefix', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
 
 DEFAULT_LOGFIRE_INSTANCE = Logfire()
 span = DEFAULT_LOGFIRE_INSTANCE.span

logfire-api/logfire_api/_internal/stack_info.pyi

@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from pathlib import Path
 from types import CodeType, FrameType
 from typing import TypedDict
 
@@ -39,3 +40,8 @@ def is_user_code(code: CodeType) -> bool:
         On the other hand, generator expressions and lambdas might be called far away from where they are defined.
     """
 def warn_at_user_stacklevel(msg: str, category: type[Warning]): ...
+def add_non_user_code_prefix(path: str | Path) -> None:
+    """Add a path to the list of prefixes that are considered non-user code.
+
+    This prevents the stack info from including frames from the given path.
+    """

logfire/__init__.py

@@ -12,6 +12,7 @@
 from ._internal.constants import LevelName
 from ._internal.main import Logfire, LogfireSpan
 from ._internal.scrubbing import ScrubbingOptions, ScrubMatch
+from ._internal.stack_info import add_non_user_code_prefix
 from ._internal.utils import suppress_instrumentation
 from .integrations.logging import LogfireLoggingHandler
 from .integrations.structlog import LogfireProcessor as StructlogProcessor
@@ -142,6 +143,7 @@ def loguru_handler() -> Any:
     'ScrubMatch',
     'ScrubbingOptions',
     'VERSION',
+    'add_non_user_code_prefix',
     'suppress_instrumentation',
     'StructlogProcessor',
     'LogfireLoggingHandler',

logfire/_internal/stack_info.py

@@ -19,10 +19,33 @@
 STACK_INFO_KEYS = set(StackInfo.__annotations__.keys())
 assert STACK_INFO_KEYS == {'code.filepath', 'code.lineno', 'code.function'}
 
-SITE_PACKAGES_DIR = str(Path(opentelemetry.sdk.trace.__file__).parent.parent.parent.parent.absolute())
-PYTHON_LIB_DIR = str(Path(inspect.__file__).parent.absolute())
-LOGFIRE_DIR = str(Path(logfire.__file__).parent.absolute())
-NON_USER_CODE_PREFIXES = (SITE_PACKAGES_DIR, PYTHON_LIB_DIR, LOGFIRE_DIR)
+NON_USER_CODE_PREFIXES: tuple[str, ...] = ()
+
+
+def add_non_user_code_prefix(path: str | Path) -> None:
+    """Add a path to the list of prefixes that are considered non-user code.
+
+    This prevents the stack info from including frames from the given path.
+
+    This is for advanced users and shouldn't often be needed.
+    By default, the following prefixes are already included:
+
+    - The standard library
+    - site-packages (specifically wherever opentelemetry is installed)
+    - The logfire package
+
+    This function is useful if you're writing a library that uses logfire and you want to exclude your library's frames.
+    Since site-packages is already included, this is already the case by default for users of your library.
+    But this is useful when testing your library since it's not installed in site-packages.
+    """
+    global NON_USER_CODE_PREFIXES
+    path = str(Path(path).absolute())
+    NON_USER_CODE_PREFIXES += (path,)  # pyright: ignore[reportConstantRedefinition]
+
+
+add_non_user_code_prefix(Path(opentelemetry.sdk.trace.__file__).parent.parent.parent.parent)
+add_non_user_code_prefix(Path(inspect.__file__).parent)
+add_non_user_code_prefix(Path(logfire.__file__).parent)
 
 
 def get_filepath_attribute(file: str) -> StackInfo:

tests/test_logfire_api.py

@@ -94,6 +94,10 @@ def test_runtime(logfire_api_factory: Callable[[], ModuleType], module_name: str
     logfire_api.no_auto_trace(lambda: None)  # pragma: no branch
     logfire__all__.remove('no_auto_trace')
 
+    assert hasattr(logfire_api, 'add_non_user_code_prefix')
+    logfire_api.add_non_user_code_prefix('/foo/bar')
+    logfire__all__.remove('add_non_user_code_prefix')
+
     assert hasattr(logfire_api, 'suppress_instrumentation')
     with logfire_api.suppress_instrumentation():
         ...

tests/test_utils.py

@@ -31,7 +31,7 @@ def test_reraise_internal_exception():
 def test_internal_exception_tb(caplog: pytest.LogCaptureFixture):
     # Pretend that `internal_logfire_code_example` is a module within logfire,
     # so all frames from it should be included.
-    logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (internal_logfire_code_example.__file__,)
+    logfire.add_non_user_code_prefix(internal_logfire_code_example.__file__)
 
     user_code_example.user1()