function decorator for adding code level attributes to span

Author: wangzlei

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/code_correlation/__init__.py

@@ -7,17 +7,13 @@
 This module provides functionality for correlating code execution with telemetry data.
 """
 
-__version__ = "1.0.0"
-
-
-"""
-Utility functions for adding code information to OpenTelemetry spans.
-"""
-
-from typing import Any, Callable
 from functools import wraps
+from typing import Any, Callable
+
 from opentelemetry import trace
 
+__version__ = "1.0.0"
+
 
 # Code correlation attribute constants
 CODE_FUNCTION_NAME = "code.function.name"
@@ -28,25 +24,25 @@
 def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
     """
     Add code-related attributes to a span based on a Python function.
-    
+
     This utility method extracts function metadata and adds the following
     span attributes:
     - CODE_FUNCTION_NAME: The name of the function
     - CODE_FILE_PATH: The file path where the function is defined
     - CODE_LINE_NUMBER: The line number where the function is defined
-    
+
     Args:
         span: The OpenTelemetry span to add attributes to
         func: The Python function to extract metadata from
     """
     if not span.is_recording():
         return
-    
+
     try:
         # Get function name
         function_name = getattr(func, '__name__', str(func))
         span.set_attribute(CODE_FUNCTION_NAME, function_name)
-        
+
         # Get function source file from code object
         try:
             if hasattr(func, '__code__'):
@@ -56,7 +52,7 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
             # Handle cases where code object is not available
             # (e.g., built-in functions, C extensions)
             pass
-            
+
         # Get function line number from code object
         try:
             if hasattr(func, '__code__'):
@@ -65,7 +61,7 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
         except (AttributeError, TypeError):
             # Handle cases where code object is not available
             pass
-            
+
     except Exception:
         # Silently handle any unexpected errors to avoid breaking
         # the instrumentation flow
@@ -75,37 +71,37 @@ def _add_code_attributes_to_span(span, func: Callable[..., Any]) -> None:
 def add_code_attributes_to_span(func: Callable[..., Any]) -> Callable[..., Any]:
     """
     Decorator to automatically add code attributes to the current OpenTelemetry span.
-    
+
     This decorator extracts metadata from the decorated function and adds it as
     attributes to the current active span. The attributes added are:
     - code.function.name: The name of the function
     - code.file.path: The file path where the function is defined
     - code.line.number: The line number where the function is defined
-    
+
     This decorator supports both synchronous and asynchronous functions.
-    
+
     Usage:
         @add_code_attributes_to_span
         def my_sync_function():
             # Sync function implementation
             pass
-            
+
         @add_code_attributes_to_span
         async def my_async_function():
             # Async function implementation
             pass
-    
+
     Args:
         func: The function to be decorated
-        
+
     Returns:
         The wrapped function with current span code attributes tracing
     """
     # Detect async functions: check function code object flags or special attributes
     # CO_ITERABLE_COROUTINE = 0x80, async functions will have this flag set
-    is_async = (hasattr(func, '__code__') and 
+    is_async = (hasattr(func, '__code__') and
                 func.__code__.co_flags & 0x80) or hasattr(func, '_is_coroutine')
-    
+
     if is_async:
         # Async function wrapper
         @wraps(func)
@@ -118,10 +114,10 @@ async def async_wrapper(*args, **kwargs):
             except Exception:
                 # Silently handle any unexpected errors
                 pass
-            
+
             # Call and await the original async function
             return await func(*args, **kwargs)
-        
+
         return async_wrapper
     else:
         # Sync function wrapper
@@ -135,8 +131,8 @@ def sync_wrapper(*args, **kwargs):
             except Exception:
                 # Silently handle any unexpected errors
                 pass
-            
+
             # Call the original sync function
             return func(*args, **kwargs)
-        
+
         return sync_wrapper

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/__init__.py

@@ -0,0 +1,2 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_code_correlation.py renamed to aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/code_correlation/test_code_correlation.py

@@ -42,11 +42,11 @@ def test_function():
 
         # Verify function name attribute is set
         self.mock_span.set_attribute.assert_any_call(CODE_FUNCTION_NAME, "test_function")
-        
+
         # Verify file path attribute is set
         expected_file_path = test_function.__code__.co_filename
         self.mock_span.set_attribute.assert_any_call(CODE_FILE_PATH, expected_file_path)
-        
+
         # Verify line number attribute is set
         expected_line_number = test_function.__code__.co_firstlineno
         self.mock_span.set_attribute.assert_any_call(CODE_LINE_NUMBER, expected_line_number)
@@ -136,7 +136,7 @@ def test_add_code_attributes_co_filename_exception(self):
         mock_func.__name__ = "test_func"
         mock_code = MagicMock()
         mock_code.co_firstlineno = 10
-        
+
         # Make co_filename raise AttributeError
         type(mock_code).co_filename = PropertyMock(side_effect=AttributeError("Test exception"))
         mock_func.__code__ = mock_code
@@ -158,7 +158,7 @@ def test_add_code_attributes_co_firstlineno_exception(self):
         mock_func.__name__ = "test_func"
         mock_code = MagicMock()
         mock_code.co_filename = "/test/file.py"
-        
+
         # Make co_firstlineno raise TypeError
         type(mock_code).co_firstlineno = PropertyMock(side_effect=TypeError("Test exception"))
         mock_func.__code__ = mock_code
@@ -180,7 +180,7 @@ def test_add_code_attributes_co_filename_type_error(self):
         mock_func.__name__ = "test_func"
         mock_code = MagicMock()
         mock_code.co_firstlineno = 10
-        
+
         # Make co_filename raise TypeError
         type(mock_code).co_filename = PropertyMock(side_effect=TypeError("Test exception"))
         mock_func.__code__ = mock_code
@@ -364,7 +364,7 @@ def test_function():
         # Verify the function still works correctly despite internal exception
         self.assertEqual(result, "test result")
 
-    @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span')  
+    @patch('amazon.opentelemetry.distro.code_correlation._add_code_attributes_to_span')
     @patch('amazon.opentelemetry.distro.code_correlation.trace.get_current_span')
     def test_decorator_internal_exception_handling_async(self, mock_get_current_span, mock_add_attributes):
         """Test that decorator handles internal exceptions gracefully in async function."""