Coverage for faststream / opentelemetry / baggage.py: 98%

1from typing import TYPE_CHECKING, Any, cast (empty)

2 

3from opentelemetry import baggage, context (empty)

4from opentelemetry.baggage.propagation import W3CBaggagePropagator (empty)

5from typing_extensions import Self (empty)

6 

7if TYPE_CHECKING: (empty)

8 from faststream.message import StreamMessage 

9 

10_BAGGAGE_PROPAGATOR = W3CBaggagePropagator() (empty)

11 

12 

13class Baggage: (empty)

14 __slots__ = ("_baggage", "_batch_baggage") (empty)

15 

16 def __init__( (empty)

17 self, 

18 payload: dict[str, Any], 

19 batch_payload: list[dict[str, Any]] | None = None, 

20 ) -> None: 

21 self._baggage = dict(payload) (empty)

22 self._batch_baggage = [dict(b) for b in batch_payload] if batch_payload else [] (empty)

23 

24 def get_all(self) -> dict[str, Any]: (empty)

25 """Get a copy of the current baggage.""" 

26 return self._baggage.copy() (empty)

27 

28 def get_all_batch(self) -> list[dict[str, Any]]: (empty)

29 """Get a copy of all batch baggage if exists.""" 

30 return self._batch_baggage.copy() (empty)

31 

32 def get(self, key: str) -> Any | None: (empty)

33 """Get a value from the baggage by key.""" 

34 return self._baggage.get(key) (empty)

35 

36 def remove(self, key: str) -> None: (empty)

37 """Remove a baggage item by key.""" 

38 self._baggage.pop(key, None) (empty)

39 

40 def set(self, key: str, value: Any) -> None: (empty)

41 """Set a key-value pair in the baggage.""" 

42 self._baggage[key] = value (empty)

43 

44 def clear(self) -> None: (empty)

45 """Clear the current baggage.""" 

46 self._baggage.clear() (empty)

47 

48 def to_headers(self, headers: dict[str, Any] | None = None) -> dict[str, Any]: (empty)

49 """Convert baggage items to headers format for propagation.""" 

50 current_context = context.get_current() (empty)

51 if headers is None: 51 ↛ 54line 51 didn't jump to line 54 because the condition on line 51 was always true(empty)

52 headers = {} (empty)

53 

54 for k, v in self._baggage.items(): (empty)

55 current_context = baggage.set_baggage(k, v, context=current_context) (empty)

56 

57 _BAGGAGE_PROPAGATOR.inject(headers, current_context) (empty)

58 return headers (empty)

59 

60 @classmethod (empty)

61 def from_msg(cls, msg: "StreamMessage[Any]") -> Self: (empty)

62 """Create a Baggage instance from a StreamMessage.""" 

63 if len(msg.batch_headers) <= 1: (empty)

64 return cls.from_headers(msg.headers) (empty)

65 

66 cumulative_baggage: dict[str, Any] = {} (empty)

67 batch_baggage: list[dict[str, Any]] = [] (empty)

68 

69 for headers in msg.batch_headers: (empty)

70 payload = baggage.get_all(_BAGGAGE_PROPAGATOR.extract(headers)) (empty)

71 cumulative_baggage.update(payload) (empty)

72 batch_baggage.append(cast("dict[str, Any]", payload)) (empty)

73 

74 return cls(cumulative_baggage, batch_baggage) (empty)

75 

76 @classmethod (empty)

77 def from_headers(cls, headers: dict[str, Any]) -> Self: (empty)

78 """Create a Baggage instance from headers.""" 

79 payload = baggage.get_all(_BAGGAGE_PROPAGATOR.extract(headers)) (empty)

80 return cls(cast("dict[str, Any]", payload)) (empty)

81 

82 def __repr__(self) -> str: (empty)

83 return self._baggage.__repr__() (empty)