Allow new context argument of metric instrument methods to be passed positionally (#616)

Author: alexmojaki

logfire/_internal/metrics.py

@@ -3,10 +3,9 @@
 import dataclasses
 from abc import ABC, abstractmethod
 from threading import Lock
-from typing import Any, Generic, Sequence, TypedDict, TypeVar
+from typing import Any, Generic, Sequence, TypeVar
 from weakref import WeakSet
 
-from opentelemetry.context import Context
 from opentelemetry.metrics import (
     CallbackT,
     Counter,
@@ -21,7 +20,6 @@
 )
 from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider
 from opentelemetry.util.types import Attributes
-from typing_extensions import Unpack
 
 try:
     # This only exists in opentelemetry-sdk>=1.23.0
@@ -213,17 +211,6 @@ def create_observable_up_down_counter(
 InstrumentT = TypeVar('InstrumentT', bound=Instrument)
 
 
-class MaybeContext(TypedDict, total=False):
-    """Backward-compatible keyword arguments for methods like `Counter.add`.
-
-    Starting with opentelemetry-sdk 1.28.0, these methods accept an additional optional `context` argument.
-    This is passed to the underlying instrument using `**kwargs` for compatibility with older versions.
-    This is the type hint for those kwargs.
-    """
-
-    context: Context | None
-
-
 class _ProxyInstrument(ABC, Generic[InstrumentT]):
     def __init__(
         self,
@@ -263,26 +250,30 @@ def __init__(
 
 
 class _ProxyCounter(_ProxyInstrument[Counter], Counter):
-    def add(  # type: ignore
+    def add(
         self,
         amount: int | float,
         attributes: Attributes | None = None,
-        **kwargs: Unpack[MaybeContext],
+        # Starting with opentelemetry-sdk 1.28.0, these methods accept an additional optional `context` argument.
+        # This is passed to the underlying instrument using `*args, **kwargs` for compatibility with older versions.
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        self._instrument.add(amount, attributes, **kwargs)
+        self._instrument.add(amount, attributes, *args, **kwargs)
 
     def _create_real_instrument(self, meter: Meter) -> Counter:
         return meter.create_counter(self._name, self._unit, self._description)
 
 
 class _ProxyHistogram(_ProxyInstrument[Histogram], Histogram):
-    def record(  # type: ignore
+    def record(
         self,
         amount: int | float,
         attributes: Attributes | None = None,
-        **kwargs: Unpack[MaybeContext],
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        self._instrument.record(amount, attributes, **kwargs)
+        self._instrument.record(amount, attributes, *args, **kwargs)
 
     def _create_real_instrument(self, meter: Meter) -> Histogram:
         return meter.create_histogram(self._name, self._unit, self._description)
@@ -310,13 +301,14 @@ def _create_real_instrument(self, meter: Meter) -> ObservableUpDownCounter:  # p
 
 
 class _ProxyUpDownCounter(_ProxyInstrument[UpDownCounter], UpDownCounter):
-    def add(  # type: ignore
+    def add(
         self,
         amount: int | float,
         attributes: Attributes | None = None,
-        **kwargs: Unpack[MaybeContext],
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        self._instrument.add(amount, attributes, **kwargs)
+        self._instrument.add(amount, attributes, *args, **kwargs)
 
     def _create_real_instrument(self, meter: Meter) -> UpDownCounter:
         return meter.create_up_down_counter(self._name, self._unit, self._description)
@@ -325,13 +317,14 @@ def _create_real_instrument(self, meter: Meter) -> UpDownCounter:
 if Gauge is not None:  # pragma: no branch
 
     class _ProxyGauge(_ProxyInstrument[Gauge], Gauge):
-        def set(  # type: ignore
+        def set(
             self,
             amount: int | float,
             attributes: Attributes | None = None,
-            **kwargs: Unpack[MaybeContext],
+            *args: Any,
+            **kwargs: Any,
         ) -> None:  # pragma: no cover
-            self._instrument.set(amount, attributes, **kwargs)
+            self._instrument.set(amount, attributes, *args, **kwargs)
 
         def _create_real_instrument(self, meter: Meter):  # pragma: no cover
             return meter.create_gauge(self._name, self._unit, self._description)

tests/test_metrics.py

@@ -7,6 +7,7 @@
 import requests
 from dirty_equals import IsInt
 from inline_snapshot import snapshot
+from opentelemetry import metrics
 from opentelemetry.metrics import CallbackOptions, Observation
 from opentelemetry.sdk.metrics._internal.export import MetricExporter, MetricExportResult
 from opentelemetry.sdk.metrics.export import AggregationTemporality, InMemoryMetricReader, MetricsData
@@ -15,6 +16,41 @@
 import logfire._internal.metrics
 from logfire._internal.exporters.quiet_metrics import QuietMetricExporter
 
+meter = metrics.get_meter('global_test_meter')
+
+global_test_counter = meter.create_counter(name='global_test_counter')
+
+
+def test_global_test_counter(metrics_reader: InMemoryMetricReader) -> None:
+    global_test_counter.add(1)
+    global_test_counter.add(20)
+    metrics_reader.collect()
+    global_test_counter.add(300)
+    global_test_counter.add(4000)
+
+    assert get_collected_metrics(metrics_reader) == snapshot(
+        [
+            {
+                'name': 'global_test_counter',
+                'description': '',
+                'unit': '',
+                'data': {
+                    'data_points': [
+                        {
+                            'attributes': {},
+                            'start_time_unix_nano': IsInt(),
+                            'time_unix_nano': IsInt(),
+                            'value': 300 + 4000,
+                            'exemplars': [],
+                        }
+                    ],
+                    'aggregation_temporality': AggregationTemporality.DELTA,
+                    'is_monotonic': True,
+                },
+            }
+        ]
+    )
+
 
 def test_create_metric_counter(metrics_reader: InMemoryMetricReader) -> None:
     counter = logfire.metric_counter('counter')