Spans and Tracing

def span(msg_template: str, /, *,
         _tags: Sequence[str] | None = None,
         _span_name: str | None = None,
         _level: LevelName | int | None = None,
         _links: Sequence = (),
         **attributes) -> LogfireSpan:
    """
    Create a span context manager for tracking an operation.
    
    Parameters:
    - msg_template: Message template with {} placeholders for attributes
    - _tags: Optional sequence of tags to apply to this span
    - _span_name: Custom span name (defaults to msg_template)
    - _level: Log level for the span
    - _links: Sequence of links to other spans
    - **attributes: Key-value attributes for structured data and template substitution
    
    Returns: LogfireSpan context manager
    """

import logfire

# Basic span usage
with logfire.span('Processing user data', user_id=123):
    # Work happens here, all logs are associated with this span
    logfire.info('Validating user input')
    # ... processing logic
    logfire.info('Processing complete')

# Nested spans for hierarchical tracking
with logfire.span('HTTP Request', method='POST', endpoint='/users'):
    with logfire.span('Authentication', user_id=123):
        # Auth logic
        pass
    
    with logfire.span('Database Query', table='users'):
        # DB query logic
        pass
    
    with logfire.span('Response Generation'):
        # Response building
        pass

# Custom span name and level
with logfire.span('Calculating metrics for {user_id}', 
                 _span_name='metric_calculation',
                 _level='debug',
                 user_id=456, calculation_type='advanced'):
    # Calculation logic
    pass

class LogfireSpan:
    """
    Context manager for spans with additional functionality.
    """
    
    @property
    def message_template(self) -> str:
        """The message template for this span (read-only)."""
    
    @property 
    def tags(self) -> tuple[str, ...]:
        """Tuple of tags for this span (read/write)."""
    
    @tags.setter
    def tags(self, value: Sequence[str]) -> None: ...
    
    @property
    def message(self) -> str:
        """The formatted message for this span (read/write)."""
    
    @message.setter
    def message(self, value: str) -> None: ...
    
    def set_attribute(self, key: str, value: Any) -> None:
        """
        Set an attribute on the span.
        
        Parameters:
        - key: Attribute key
        - value: Attribute value (will be converted to string if needed)
        """
    
    def set_attributes(self, attributes: dict[str, Any]) -> None:
        """
        Set multiple attributes on the span.
        
        Parameters:
        - attributes: Dictionary of key-value attribute pairs
        """
    
    def add_link(self, context: SpanContext, attributes: dict[str, Any] | None = None) -> None:
        """
        Add a link to another span.
        
        Parameters:
        - context: SpanContext of the span to link to
        - attributes: Optional attributes for the link
        """
    
    def record_exception(self, exception: BaseException, *,
                        attributes: dict[str, Any] | None = None,
                        timestamp: int | None = None,
                        escaped: bool = False) -> None:
        """
        Record an exception on the span.
        
        Parameters:
        - exception: Exception instance to record
        - attributes: Optional attributes associated with the exception
        - timestamp: Optional timestamp (nanoseconds since epoch)
        - escaped: Whether the exception was handled/escaped
        """
    
    def is_recording(self) -> bool:
        """
        Check if the span is currently recording.
        
        Returns: True if the span is recording, False otherwise
        """
    
    def set_level(self, level: LevelName | int) -> None:
        """
        Set the log level of the span.
        
        Parameters:
        - level: Log level name or integer value
        """

import logfire

# Dynamic span modification
with logfire.span('Processing request') as span:
    span.set_attribute('start_time', time.time())
    
    try:
        # Processing logic
        result = process_data()
        span.set_attributes({
            'result_count': len(result),
            'success': True
        })
    except Exception as e:
        span.record_exception(e, escaped=True)
        span.set_attribute('error_handled', True)
        # Handle error
    
    # Update span message
    span.message = f'Processed {len(result)} items'

# Linking spans across services
with logfire.span('External API call') as span:
    # Get context from current span to link from remote service
    current_context = span.get_span_context()
    
    # In remote service, create linked span
    with logfire.span('Remote processing') as remote_span:
        remote_span.add_link(current_context, {'service': 'external-api'})

def instrument(msg_template: str | None = None, *,
               span_name: str | None = None,
               extract_args: bool = True,
               record_return: bool = False,
               allow_generator: bool = False):
    """
    Decorator to instrument functions with automatic span creation.
    
    Parameters:
    - msg_template: Optional message template for the span
    - span_name: Optional custom span name
    - extract_args: Whether to automatically log function arguments as attributes
    - record_return: Whether to log the return value as an attribute
    - allow_generator: Allow instrumenting generator functions
    
    Returns: Decorator function
    """

import logfire

# Basic function instrumentation
@logfire.instrument
def calculate_total(items):
    return sum(item.price for item in items)

# Custom message and span name
@logfire.instrument(
    msg_template='Computing metrics for {period}',
    span_name='metric_computation'
)
def compute_metrics(period, data_source):
    # Computation logic
    return {'avg': 10, 'count': 100}

# Record return values
@logfire.instrument(record_return=True)
def fetch_user_data(user_id):
    # Returns user data that will be logged as span attribute
    return {'id': user_id, 'name': 'John', 'email': 'john@example.com'}

# Instrument without argument extraction (for sensitive data)
@logfire.instrument(extract_args=False)
def process_payment(credit_card_info):
    # Credit card info won't be logged automatically
    return payment_result

# Generator function instrumentation  
@logfire.instrument(allow_generator=True)
def process_large_dataset():
    for batch in large_dataset:
        yield process_batch(batch)

def install_auto_tracing(modules: Sequence[str] | Callable[[AutoTraceModule], bool], *,
                        min_duration: float,
                        check_imported_modules: Literal['error', 'warn', 'ignore'] = 'error') -> None:
    """
    Install automatic tracing for specified modules.
    
    Parameters:
    - modules: Module names to trace, or a predicate function
    - min_duration: Minimum execution duration (seconds) to create spans
    - check_imported_modules: How to handle already-imported modules
    """

def no_auto_trace(func):
    """
    Decorator to exclude a function from auto-tracing.
    
    Parameters:
    - func: Function to exclude from auto-tracing
    
    Returns: Original function with auto-trace exclusion marker
    """

import logfire

# Auto-trace specific modules
logfire.install_auto_tracing(
    modules=['myapp.services', 'myapp.models'],
    min_duration=0.01  # Only trace functions taking > 10ms
)

# Auto-trace with predicate function using AutoTraceModule
def should_trace(module: AutoTraceModule):
    return (module.name.startswith('myapp.') and 
            not module.name.endswith('.tests') and
            module.filename is not None)

logfire.install_auto_tracing(
    modules=should_trace,
    min_duration=0.005
)

# Exclude specific functions from auto-tracing
@logfire.no_auto_trace
def internal_helper_function():
    # This won't be auto-traced even if module is included
    pass

def log_slow_async_callbacks(slow_duration: float = 0.1) -> AbstractContextManager[None]:
    """
    Context manager that logs warnings for slow asyncio callbacks.
    
    Parameters:
    - slow_duration: Threshold in seconds for considering callbacks slow
    
    Returns: Context manager for monitoring async callback performance
    """

import asyncio
import logfire

# Async spans work automatically
async def async_operation():
    with logfire.span('Async database query', query_type='SELECT'):
        await asyncio.sleep(0.1)  # Simulated async work
        return {'results': []}

# Monitor slow async callbacks
async def main():
    with logfire.log_slow_async_callbacks(slow_duration=0.05):
        # Any callback taking > 50ms will be logged as warning
        await async_operation()
        await another_async_operation()

# Async function instrumentation
@logfire.instrument
async def fetch_user_async(user_id):
    with logfire.span('Database lookup', user_id=user_id):
        # Async database call
        return await db.fetch_user(user_id)

def suppress_scopes(*scopes: str) -> None:
    """
    Prevent span and metric creation for specified OpenTelemetry scopes.
    
    Parameters:
    - *scopes: OpenTelemetry scope names to suppress
    """

import logfire

# Suppress spans from specific instrumentation
logfire.suppress_scopes('opentelemetry.instrumentation.requests')

# Now requests won't create spans, but other instrumentation still works
import requests
requests.get('https://api.example.com')  # No span created

# Custom instrumentation still works
with logfire.span('Custom operation'):
    requests.get('https://api.example.com')  # Still no request span, but custom span exists

def get_baggage() -> dict[str, str]:
    """
    Get current OpenTelemetry baggage (key-value pairs propagated across service boundaries).
    
    Returns: Dictionary of baggage key-value pairs
    """

def set_baggage(baggage: dict[str, str]) -> Token:
    """
    Set OpenTelemetry baggage for context propagation.
    
    Parameters:
    - baggage: Dictionary of key-value pairs to propagate
    
    Returns: Token for context restoration
    """

import logfire

# Set baggage for cross-service correlation
token = logfire.set_baggage({
    'user_id': '123',
    'session_id': 'abc-def-456',
    'feature_flag': 'new_checkout_enabled'
})

# Baggage is automatically propagated in HTTP headers
with logfire.span('Service A operation'):
    # Make HTTP call - baggage is automatically included in headers
    response = requests.post('https://service-b/process')

# In Service B, retrieve baggage
current_baggage = logfire.get_baggage()
user_id = current_baggage.get('user_id')
feature_enabled = current_baggage.get('feature_flag') == 'new_checkout_enabled'

def force_flush(timeout_millis: int = 3000) -> bool:
    """
    Force flush all pending spans to configured exporters.
    
    Parameters:
    - timeout_millis: Maximum time to wait for flush completion
    
    Returns: True if flush completed within timeout
    """

def shutdown(timeout_millis: int = 30000, flush: bool = True) -> bool:
    """
    Shutdown span processors and exporters.
    
    Parameters:
    - timeout_millis: Maximum time to wait for shutdown
    - flush: Whether to flush pending spans before shutdown
    
    Returns: True if shutdown completed within timeout
    """

import logfire
import atexit

# Ensure spans are flushed on application exit
def cleanup():
    logfire.force_flush(timeout_millis=5000)
    logfire.shutdown(timeout_millis=10000)

atexit.register(cleanup)

# Manual flush at critical points
def handle_request():
    with logfire.span('Handle request'):
        # Process request
        pass
    
    # Ensure this request's spans are sent before continuing
    logfire.force_flush(timeout_millis=1000)

# OpenTelemetry types used in span operations
from opentelemetry.trace import SpanContext, SpanKind
from opentelemetry.util.types import AttributeValue

# Logfire-specific types
LevelName = Literal['trace', 'debug', 'info', 'notice', 'warn', 'warning', 'error', 'fatal']

# Auto-tracing types
@dataclass
class AutoTraceModule:
    """
    Information about a module being imported that should maybe be traced automatically.
    
    This object will be passed to a function that should return True if the module should be traced.
    Used with install_auto_tracing() as the modules argument.
    """
    
    name: str
    """Fully qualified absolute name of the module being imported."""
    
    filename: str | None
    """Filename of the module being imported."""
    
    def parts_start_with(self, prefix: str | Sequence[str]) -> bool:
        """
        Return True if the module name starts with any of the given prefixes, using dots as boundaries.
        
        For example, if the module name is 'foo.bar.spam', then parts_start_with('foo') will return True,
        but parts_start_with('bar') or parts_start_with('foo_bar') will return False.
        
        If a prefix contains any characters other than letters, numbers, and dots,
        then it will be treated as a regular expression.
        
        Parameters:
        - prefix: String or sequence of strings to match against module name
        
        Returns: True if module name matches any prefix
        """

# Context management
from contextvars import Token
from typing import AbstractContextManager