spanwire 0.1.0__py3-none-any.whl

spanwire/core.py

@@ -0,0 +1,463 @@
+"""Core primitives of spanwire: SpanContext, Span, Tracer, and SpanStatus."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import inspect
+import secrets
+import time
+import traceback
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from spanwire.breadcrumbs import Breadcrumb, BreadcrumbBuffer
+from spanwire.context import (
+    clear_baggage,
+    get_baggage,
+    get_current_span,
+    reset_current_span,
+    set_baggage,
+    set_current_span,
+)
+from spanwire.exporters import InMemoryExporter, SpanExporter
+from spanwire.sampler import AlwaysOnSampler, Sampler
+from spanwire.w3c import (
+    extract_w3c_headers,
+    format_traceparent,
+    inject_w3c_headers,
+    is_valid_span_id,
+    is_valid_trace_id,
+)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def generate_trace_id() -> str:
+    """Generate a random 128-bit (32 hex characters) trace ID."""
+    return secrets.token_hex(16)
+
+
+def generate_span_id() -> str:
+    """Generate a random 64-bit (16 hex characters) span ID."""
+    return secrets.token_hex(8)
+
+
+class SpanStatus(str, Enum):
+    """Execution status of a span."""
+
+    UNSET = "UNSET"
+    OK = "OK"
+    ERROR = "ERROR"
+
+
+@dataclass(frozen=True)
+class SpanContext:
+    """Immutable context identifying a span across distributed components."""
+
+    trace_id: str
+    span_id: str
+    trace_flags: int = 1
+    tracestate: str = ""
+
+    def __post_init__(self) -> None:
+        if not is_valid_trace_id(self.trace_id):
+            raise ValueError(f"Invalid trace_id: {self.trace_id}")
+        if not is_valid_span_id(self.span_id):
+            raise ValueError(f"Invalid span_id: {self.span_id}")
+
+    @property
+    def is_sampled(self) -> bool:
+        """Return True if the sampled flag is set."""
+        return bool(self.trace_flags & 1)
+
+    def to_traceparent(self) -> str:
+        """Format as W3C traceparent string."""
+        return format_traceparent(self.trace_id, self.span_id, self.trace_flags)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert context to dictionary."""
+        return {
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+            "trace_flags": self.trace_flags,
+            "tracestate": self.tracestate,
+            "is_sampled": self.is_sampled,
+        }
+
+
+@dataclass
+class SpanEvent:
+    """A timed event attached to a span with arbitrary attributes."""
+
+    name: str
+    timestamp_ns: int = field(default_factory=time.time_ns)
+    attributes: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "name": self.name,
+            "timestamp_ns": self.timestamp_ns,
+            "attributes": dict(self.attributes),
+        }
+
+
+class Span:
+    """Represents a single timed unit of work in a trace."""
+
+    def __init__(
+        self,
+        name: str,
+        context: SpanContext,
+        parent_span_id: Optional[str] = None,
+        attributes: Optional[Mapping[str, Any]] = None,
+        tracer: Optional[Tracer] = None,
+    ) -> None:
+        self.name = name
+        self.context = context
+        self.parent_span_id = parent_span_id
+        self._tracer = tracer
+
+        self.start_time_ns: int = time.time_ns()
+        self.end_time_ns: Optional[int] = None
+        self.status: SpanStatus = SpanStatus.UNSET
+        self.status_description: str = ""
+
+        self.attributes: Dict[str, Any] = dict(attributes) if attributes else {}
+        self.events: List[SpanEvent] = []
+        self.breadcrumbs: List[Breadcrumb] = []
+        self._token: Optional[Any] = None
+        self._ended: bool = False
+
+    @property
+    def duration_ns(self) -> int:
+        """Duration in nanoseconds."""
+        end = self.end_time_ns if self.end_time_ns is not None else time.time_ns()
+        return max(0, end - self.start_time_ns)
+
+    @property
+    def duration_ms(self) -> float:
+        """Duration in milliseconds."""
+        return self.duration_ns / 1_000_000.0
+
+    @property
+    def is_ended(self) -> bool:
+        return self._ended
+
+    def set_attribute(self, key: str, value: Any) -> Span:
+        """Set a single key-value attribute."""
+        if not self._ended:
+            self.attributes[str(key)] = value
+        return self
+
+    def set_attributes(self, mapping: Mapping[str, Any]) -> Span:
+        """Set multiple attributes at once."""
+        if not self._ended and mapping:
+            for k, v in mapping.items():
+                self.attributes[str(k)] = v
+        return self
+
+    def add_event(
+        self,
+        name: str,
+        attributes: Optional[Mapping[str, Any]] = None,
+    ) -> SpanEvent:
+        """Add a timed milestone event to this span."""
+        event = SpanEvent(
+            name=name,
+            timestamp_ns=time.time_ns(),
+            attributes=dict(attributes) if attributes else {},
+        )
+        if not self._ended:
+            self.events.append(event)
+        return event
+
+    def set_status(
+        self,
+        status: Union[SpanStatus, str],
+        description: str = "",
+    ) -> Span:
+        """Set the completion status of the span."""
+        if not self._ended:
+            if isinstance(status, str):
+                self.status = SpanStatus(status.upper())
+            else:
+                self.status = status
+            self.status_description = description
+        return self
+
+    def record_exception(
+        self,
+        exc: BaseException,
+        escaped: bool = True,
+    ) -> Span:
+        """Record an exception on this span, setting status to ERROR."""
+        self.set_status(SpanStatus.ERROR, str(exc))
+        tb = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
+        self.add_event(
+            "exception",
+            {
+                "exception.type": type(exc).__name__,
+                "exception.message": str(exc),
+                "exception.stacktrace": tb,
+                "exception.escaped": escaped,
+            },
+        )
+        return self
+
+    def add_breadcrumb(
+        self,
+        message: str,
+        level: str = "INFO",
+        category: str = "default",
+        **data: Any,
+    ) -> Breadcrumb:
+        """Record a breadcrumb specifically attached to this span."""
+        crumb = Breadcrumb(
+            message=message,
+            level=level.upper(),
+            category=category,
+            timestamp=time.time(),
+            data=data,
+        )
+        self.breadcrumbs.append(crumb)
+        return crumb
+
+    def end(self) -> None:
+        """Mark the span as completed and dispatch to exporter."""
+        if self._ended:
+            return
+        self.end_time_ns = time.time_ns()
+        self._ended = True
+
+        if self.status == SpanStatus.UNSET:
+            self.status = SpanStatus.OK
+
+        # Collect breadcrumbs from tracer buffer if available
+        if self._tracer and hasattr(self._tracer, "breadcrumbs"):
+            recent = self._tracer.breadcrumbs.get_crumbs()
+            if recent and not self.breadcrumbs:
+                self.breadcrumbs.extend(recent)
+
+        if self._tracer:
+            self._tracer._on_span_ended(self)
+
+    def __enter__(self) -> Span:
+        self._token = set_current_span(self)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[Any],
+    ) -> bool:
+        try:
+            if exc_val is not None:
+                self.record_exception(exc_val)
+            self.end()
+        finally:
+            if self._token is not None:
+                reset_current_span(self._token)
+                self._token = None
+        return False  # Do not suppress exceptions
+
+    async def __aenter__(self) -> Span:
+        return self.__enter__()
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[Any],
+    ) -> bool:
+        return self.__exit__(exc_type, exc_val, exc_tb)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize span data to dictionary for JSON export."""
+        return {
+            "name": self.name,
+            "context": self.context.to_dict(),
+            "parent_span_id": self.parent_span_id,
+            "start_time_ns": self.start_time_ns,
+            "end_time_ns": self.end_time_ns,
+            "duration_ms": round(self.duration_ms, 3),
+            "status": self.status.value,
+            "status_description": self.status_description,
+            "attributes": self.attributes,
+            "events": [e.to_dict() for e in self.events],
+            "breadcrumbs": [b.to_dict() for b in self.breadcrumbs],
+        }
+
+    def __repr__(self) -> str:
+        status_name = self.status.name
+        return (
+            f"Span(name={self.name!r}, trace_id={self.context.trace_id}, "
+            f"span_id={self.context.span_id}, status={status_name})"
+        )
+
+
+class Tracer:
+    """Manages active spans, context propagation, breadcrumbs, and export."""
+
+    def __init__(
+        self,
+        service_name: str = "unnamed-service",
+        sampler: Optional[Sampler] = None,
+        exporter: Optional[SpanExporter] = None,
+        breadcrumb_capacity: int = 50,
+    ) -> None:
+        self.service_name = service_name
+        self.sampler: Sampler = sampler if sampler is not None else AlwaysOnSampler()
+        self.exporter: Optional[SpanExporter] = exporter
+        self.breadcrumbs = BreadcrumbBuffer(capacity=breadcrumb_capacity)
+
+    def current_span(self) -> Optional[Span]:
+        """Return the currently active span in context."""
+        return get_current_span()
+
+    def current_context(self) -> Optional[SpanContext]:
+        """Return the SpanContext of the currently active span, if any."""
+        span = self.current_span()
+        return span.context if span else None
+
+    def span(
+        self,
+        name: str,
+        attributes: Optional[Mapping[str, Any]] = None,
+        parent_context: Optional[SpanContext] = None,
+    ) -> Span:
+        """Create and start a new Span as a child of the active context or parent_context."""
+        parent = parent_context or self.current_context()
+
+        if parent:
+            trace_id = parent.trace_id
+            parent_span_id = parent.span_id
+            tracestate = parent.tracestate
+        else:
+            trace_id = generate_trace_id()
+            parent_span_id = None
+            tracestate = ""
+
+        span_id = generate_span_id()
+        is_sampled = self.sampler.should_sample(parent, trace_id, name)
+        trace_flags = 1 if is_sampled else 0
+
+        ctx = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            trace_flags=trace_flags,
+            tracestate=tracestate,
+        )
+
+        merged_attrs = {"service.name": self.service_name}
+        if attributes:
+            merged_attrs.update(attributes)
+
+        return Span(
+            name=name,
+            context=ctx,
+            parent_span_id=parent_span_id,
+            attributes=merged_attrs,
+            tracer=self,
+        )
+
+    def add_breadcrumb(
+        self,
+        message: str,
+        level: str = "INFO",
+        category: str = "default",
+        data: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Breadcrumb:
+        """Record a breadcrumb in the tracer buffer."""
+        return self.breadcrumbs.record(
+            message=message,
+            level=level,
+            category=category,
+            data=data,
+            **kwargs,
+        )
+
+    def inject_headers(
+        self,
+        headers: Optional[Dict[str, str]] = None,
+        context: Optional[SpanContext] = None,
+    ) -> Dict[str, str]:
+        """Inject current span's trace context and baggage into HTTP headers."""
+        ctx = context or self.current_context()
+        if not ctx:
+            return headers if headers is not None else {}
+        baggage = get_baggage()
+        return inject_w3c_headers(
+            trace_id=ctx.trace_id,
+            span_id=ctx.span_id,
+            trace_flags=ctx.trace_flags,
+            baggage=baggage,
+            tracestate=ctx.tracestate,
+            target_headers=headers,
+        )
+
+    def extract_headers(
+        self,
+        headers: Mapping[str, str],
+    ) -> Tuple[Optional[SpanContext], Dict[str, str], str]:
+        """Extract W3C trace context, baggage, and tracestate from incoming headers."""
+        tp_parsed, baggage, tracestate = extract_w3c_headers(headers)
+        if tp_parsed is None:
+            return None, baggage, tracestate
+
+        trace_id, parent_id, flags = tp_parsed
+        ctx = SpanContext(
+            trace_id=trace_id,
+            span_id=parent_id,
+            trace_flags=flags,
+            tracestate=tracestate,
+        )
+        return ctx, baggage, tracestate
+
+    def _on_span_ended(self, span: Span) -> None:
+        """Called when a span completes."""
+        if self.exporter is not None and span.context.is_sampled:
+            self.exporter.export([span])
+
+    def trace(
+        self,
+        name: Optional[str] = None,
+        attributes: Optional[Mapping[str, Any]] = None,
+    ) -> Callable[[F], F]:
+        """Decorator for tracing both synchronous and asynchronous functions."""
+
+        def decorator(fn: F) -> F:
+            span_name = name or fn.__qualname__
+
+            if inspect.iscoroutinefunction(fn):
+
+                @functools.wraps(fn)
+                async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                    with self.span(span_name, attributes=attributes):
+                        return await fn(*args, **kwargs)
+
+                return cast(F, async_wrapper)
+            else:
+
+                @functools.wraps(fn)
+                def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                    with self.span(span_name, attributes=attributes):
+                        return fn(*args, **kwargs)
+
+                return cast(F, sync_wrapper)
+
+        return decorator

spanwire/exporters.py

@@ -0,0 +1,188 @@
+"""Exporters for dispatching finished spans to telemetry collectors, logs, or memory."""
+
+from __future__ import annotations
+
+import io
+import json
+import os
+import sys
+import threading
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, TextIO, Union
+
+if TYPE_CHECKING:
+    from spanwire.core import Span
+
+
+class SpanExporter(ABC):
+    """Abstract base class for span exporters."""
+
+    @abstractmethod
+    def export(self, spans: Sequence[Span]) -> None:
+        """Export a batch of finished spans."""
+        pass
+
+    def shutdown(self) -> None:
+        """Clean up exporter resources."""
+        pass
+
+
+class InMemoryExporter(SpanExporter):
+    """Stores exported spans in memory, ideal for test suites and assertions."""
+
+    def __init__(self) -> None:
+        self._spans: List[Span] = []
+        self._lock = threading.Lock()
+
+    def export(self, spans: Sequence[Span]) -> None:
+        with self._lock:
+            self._spans.extend(spans)
+
+    def get_finished_spans(self) -> List[Span]:
+        """Return a copy of all exported spans."""
+        with self._lock:
+            return list(self._spans)
+
+    def clear(self) -> None:
+        """Clear all stored spans."""
+        with self._lock:
+            self._spans.clear()
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._spans)
+
+    def __repr__(self) -> str:
+        return f"InMemoryExporter(count={len(self)})"
+
+
+class JsonLinesExporter(SpanExporter):
+    """Exports finished spans formatted as newline-delimited JSON (NDJSON)."""
+
+    def __init__(
+        self,
+        stream_or_path: Optional[Union[TextIO, str]] = None,
+        auto_flush: bool = True,
+    ) -> None:
+        self._lock = threading.Lock()
+        self._auto_flush = auto_flush
+        self._owned_file: Optional[TextIO] = None
+
+        if stream_or_path is None:
+            self._stream: TextIO = sys.stdout
+        elif isinstance(stream_or_path, str):
+            self._owned_file = open(stream_or_path, "a", encoding="utf-8")
+            self._stream = self._owned_file
+        else:
+            self._stream = stream_or_path
+
+    def export(self, spans: Sequence[Span]) -> None:
+        with self._lock:
+            for span in spans:
+                record = span.to_dict()
+                line = json.dumps(record, default=str)
+                self._stream.write(line + "\n")
+            if self._auto_flush:
+                self._stream.flush()
+
+    def shutdown(self) -> None:
+        with self._lock:
+            if self._auto_flush:
+                self._stream.flush()
+            if self._owned_file is not None and not self._owned_file.closed:
+                self._owned_file.close()
+
+    def __repr__(self) -> str:
+        return "JsonLinesExporter()"
+
+
+class ConsoleExporter(SpanExporter):
+    """Human-readable console output exporter with ANSI colors."""
+
+    COLOR_RESET = "\033[0m"
+    COLOR_CYAN = "\033[36m"
+    COLOR_GREEN = "\033[32m"
+    COLOR_YELLOW = "\033[33m"
+    COLOR_RED = "\033[31m"
+    COLOR_GRAY = "\033[90m"
+
+    def __init__(
+        self,
+        stream: Optional[TextIO] = None,
+        colored: Optional[bool] = None,
+    ) -> None:
+        self._stream: TextIO = stream if stream is not None else sys.stderr
+        self._lock = threading.Lock()
+
+        if colored is not None:
+            self.colored = colored
+        else:
+            # Enable colors if running in a TTY or NO_COLOR is not set
+            is_tty = hasattr(self._stream, "isatty") and self._stream.isatty()
+            self.colored = is_tty and ("NO_COLOR" not in os.environ)
+
+    def _c(self, text: str, code: str) -> str:
+        if not self.colored:
+            return text
+        return f"{code}{text}{self.COLOR_RESET}"
+
+    def export(self, spans: Sequence[Span]) -> None:
+        with self._lock:
+            for span in spans:
+                status_str = span.status.name
+                if status_str == "OK":
+                    status_badge = self._c("[OK]", self.COLOR_GREEN)
+                elif status_str == "ERROR":
+                    status_badge = self._c("[ERROR]", self.COLOR_RED)
+                else:
+                    status_badge = self._c("[UNSET]", self.COLOR_GRAY)
+
+                duration_str = self._c(f"{span.duration_ms:.2f}ms", self.COLOR_YELLOW)
+                trace_info = self._c(
+                    f"trace={span.context.trace_id[:8]}.. span={span.context.span_id}",
+                    self.COLOR_CYAN,
+                )
+                parent_info = (
+                    f" parent={span.parent_span_id}" if span.parent_span_id else ""
+                )
+
+                line = f"⚡ {status_badge} {span.name} ({duration_str}) [{trace_info}{parent_info}]"
+                self._stream.write(line + "\n")
+
+                if span.attributes:
+                    attrs_formatted = ", ".join(
+                        f"{k}={v!r}" for k, v in span.attributes.items()
+                    )
+                    self._stream.write(f"   ↳ attributes: {attrs_formatted}\n")
+
+                if span.breadcrumbs:
+                    self._stream.write(f"   ↳ breadcrumbs ({len(span.breadcrumbs)}):\n")
+                    for b in span.breadcrumbs[-5:]:
+                        b_line = f"     - [{b.level}] {b.category}: {b.message}"
+                        self._stream.write(f"{b_line}\n")
+
+                if span.status.name == "ERROR" and span.status_description:
+                    self._stream.write(
+                        self._c(f"   ↳ error: {span.status_description}\n", self.COLOR_RED)
+                    )
+
+            self._stream.flush()
+
+    def __repr__(self) -> str:
+        return f"ConsoleExporter(colored={self.colored})"
+
+
+class CallbackExporter(SpanExporter):
+    """Invokes a callable for each exported span."""
+
+    def __init__(self, callback: Callable[[Span], None]) -> None:
+        self.callback = callback
+        self._lock = threading.Lock()
+
+    def export(self, spans: Sequence[Span]) -> None:
+        with self._lock:
+            for span in spans:
+                self.callback(span)
+
+    def __repr__(self) -> str:
+        return f"CallbackExporter({self.callback!r})"