merge common code between capture_span and async_capture_span (elastic#1417)

this should avoid issues with changes only being implemented
in one of the two context managers

Author: beniwohli

elasticapm/contrib/asyncio/traces.py

@@ -29,9 +29,11 @@
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import functools
+from types import TracebackType
+from typing import Optional, Type
 
 from elasticapm.conf.constants import LABEL_RE
-from elasticapm.traces import DroppedSpan, capture_span, error_logger, execution_context
+from elasticapm.traces import SpanType, capture_span, execution_context
 from elasticapm.utils import get_name_from_func
 
 
@@ -46,34 +48,13 @@ async def decorated(*args, **kwds):
 
         return decorated
 
-    async def __aenter__(self):
-        transaction = execution_context.get_transaction()
-        if transaction and transaction.is_sampled:
-            return transaction.begin_span(
-                self.name,
-                self.type,
-                context=self.extra,
-                leaf=self.leaf,
-                labels=self.labels,
-                span_subtype=self.subtype,
-                span_action=self.action,
-                sync=False,
-                start=self.start,
-            )
+    async def __aenter__(self) -> Optional[SpanType]:
+        return self.handle_enter(False)
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        transaction = execution_context.get_transaction()
-        if transaction and transaction.is_sampled:
-            try:
-                span = transaction.end_span(self.skip_frames)
-                if exc_val and not isinstance(span, DroppedSpan):
-                    try:
-                        exc_val._elastic_apm_span_id = span.id
-                    except AttributeError:
-                        # could happen if the exception has __slots__
-                        pass
-            except LookupError:
-                error_logger.debug("ended non-existing span %s of type %s", self.name, self.type)
+    async def __aexit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ):
+        self.handle_exit(exc_type, exc_val, exc_tb)
 
 
 async def set_context(data, key="custom"):

elasticapm/traces.py

@@ -35,7 +35,8 @@
 import time
 import timeit
 from collections import defaultdict
-from typing import Any, Callable, Dict, Optional, Tuple, Union
+from types import TracebackType
+from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
 
 from elasticapm.conf import constants
 from elasticapm.conf.constants import LABEL_RE, SPAN, TRANSACTION
@@ -867,7 +868,15 @@ def decorated(*args, **kwds):
 
         return decorated
 
-    def __enter__(self) -> Union[Span, DroppedSpan, None]:
+    def __enter__(self) -> Optional[SpanType]:
+        return self.handle_enter(self.sync)
+
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        self.handle_exit(exc_type, exc_val, exc_tb)
+
+    def handle_enter(self, sync: bool) -> Optional[SpanType]:
         transaction = execution_context.get_transaction()
         if transaction and transaction.is_sampled:
             return transaction.begin_span(
@@ -879,11 +888,13 @@ def __enter__(self) -> Union[Span, DroppedSpan, None]:
                 span_subtype=self.subtype,
                 span_action=self.action,
                 start=self.start,
-                sync=self.sync,
+                sync=sync,
             )
         return None
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def handle_exit(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         transaction = execution_context.get_transaction()
 
         if transaction and transaction.is_sampled: