diff --git a/requirements.txt b/requirements.txt index a708362..971b115 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ numpy>=2.0.0 pandas>=2.2.3 scikit-learn>=1.6.1 -typing_extensions>=4.12.2 \ No newline at end of file +typing_extensions>=4.12.2 +httpx~=0.28.1 \ No newline at end of file diff --git a/tests/test_analytics_http_client.py b/tests/test_analytics_http_client.py new file mode 100644 index 0000000..2b3aac7 --- /dev/null +++ b/tests/test_analytics_http_client.py @@ -0,0 +1,78 @@ +import unittest +from unittest.mock import patch, MagicMock +import httpx +from usage_analytics import AnalyticsHttpClient, ANALYTICS_HEADER_CONFIG + + +class TestAnalyticsHttpClient(unittest.TestCase): + def setUp(self) -> None: + self.module_name = "test_module" + self.client = AnalyticsHttpClient(module_name=self.module_name) + self.local_host = "http://localhost:8000" + + def test_init(self): + """Test that the client initializes with the correct module name.""" + self.assertEqual(self.client.module_name, self.module_name) + self.assertIsInstance(self.client, httpx.Client) + + @patch("usage_analytics.analytics_http_client.super") + def test_request_adds_headers(self, mock_super: MagicMock): + """Test that request method adds analytics headers.""" + mock_request = MagicMock() + mock_super.return_value.request = mock_request + + # Call request method + self.client.request("GET", self.local_host) + + # Check that super().request was called with modified headers + args, kwargs = mock_request.call_args + self.assertEqual(args[0], "GET") + self.assertEqual(args[1], self.local_host) + + # Verify headers were added + headers = kwargs.get("headers", {}) + self._validate_headers(headers) + + @patch("usage_analytics.analytics_http_client.super") + def test_stream_adds_headers(self, mock_super): + """Test that stream method adds analytics headers.""" + mock_stream = MagicMock() + mock_super.return_value.stream = mock_stream + + # Call stream method + self.client.stream("GET", self.local_host) + + # Check that super().stream was called with modified headers + args, kwargs = mock_stream.call_args + self.assertEqual(args[0], "GET") + self.assertEqual(args[1], self.local_host) + + # Verify headers were added + headers = kwargs.get("headers", {}) + self._validate_headers(headers) + + def _validate_headers(self, headers: dict) -> None: + for header_name, _ in ANALYTICS_HEADER_CONFIG: + self.assertIn(header_name, headers) + self.assertIsNotNone(headers.get(header_name)) + + def test_set_module_name(self): + """Test that the module name can be updated.""" + new_module_name = "new_test_module" + self.client.set_module_name(new_module_name) + + with patch("usage_analytics.analytics_http_client.super") as mock_super: + mock_request = MagicMock() + mock_super.return_value.request = mock_request + + self.client.request("GET", self.local_host) + + _, kwargs = mock_request.call_args + headers = kwargs.get("headers", {}) + + self.assertEqual(headers.get("PL-Module-Name"), new_module_name) + self.assertEqual(self.client.module_name, new_module_name) + + +if __name__ == "__main__": + unittest.main() diff --git a/usage_analytics/__init__.py b/usage_analytics/__init__.py new file mode 100644 index 0000000..8b07b3e --- /dev/null +++ b/usage_analytics/__init__.py @@ -0,0 +1,7 @@ +from .analytics_http_client import AnalyticsHttpClient +from .analytics_config import ANALYTICS_HEADER_CONFIG + +__all__ = [ + "AnalyticsHttpClient", + "ANALYTICS_HEADER_CONFIG", +] diff --git a/usage_analytics/analytics_config.py b/usage_analytics/analytics_config.py new file mode 100644 index 0000000..35e9dd0 --- /dev/null +++ b/usage_analytics/analytics_config.py @@ -0,0 +1,23 @@ +from enum import Enum +from typing import Callable, Optional, Tuple + +from .analytics_func import ( + get_calling_class, + get_python_version, + get_unique_call_id, +) + + +class ANALYTICS_KEYS(str, Enum): + UniqueCallID = "PL-Unique-Call-Id" + PythonVersion = "PL-Python-Version" + CallingClass = "PL-Calling-Class" + ModuleName = "PL-Module-Name" + + +ANALYTICS_HEADER_CONFIG: Tuple[Tuple[str, Optional[Callable]], ...] = ( + (ANALYTICS_KEYS.UniqueCallID, get_unique_call_id), + (ANALYTICS_KEYS.PythonVersion, get_python_version), + (ANALYTICS_KEYS.CallingClass, get_calling_class), + (ANALYTICS_KEYS.ModuleName, None), # Value provided by AnalyticsHttpClient +) diff --git a/usage_analytics/analytics_func.py b/usage_analytics/analytics_func.py new file mode 100644 index 0000000..5aae484 --- /dev/null +++ b/usage_analytics/analytics_func.py @@ -0,0 +1,36 @@ +import inspect +import platform +import uuid + + +def get_calling_class(recursive=True) -> str: + """ + Traverses the call stack to find and return the class name of the first + caller that has a 'self' variable in its frame (i.e., a method call). + If no such frame is found, returns 'StandaloneFunction'. + """ + + # Skip the current frame and the immediate caller (i.e., start from index 1) + stack = inspect.stack()[1:] + + outermost_caller = "StandaloneFunction" # Default return value + + for frame_record in stack: + frame = frame_record.frame + # Look for 'self' in the frame's local variables + self_obj = frame.f_locals.get("self", None) + if self_obj is not None: + # Return the class name of the 'self' instance + outermost_caller = type(self_obj).__name__ + if not recursive: + break + + return outermost_caller + + +def get_python_version() -> str: + return platform.python_version() + + +def get_unique_call_id() -> str: + return str(uuid.uuid4()) diff --git a/usage_analytics/analytics_http_client.py b/usage_analytics/analytics_http_client.py new file mode 100644 index 0000000..98b99e2 --- /dev/null +++ b/usage_analytics/analytics_http_client.py @@ -0,0 +1,71 @@ +import httpx +from typing import Callable + +from .analytics_config import ANALYTICS_HEADER_CONFIG, ANALYTICS_KEYS + + +class AnalyticsHttpClient(httpx.Client): + """ + Custom HTTP client that extends httpx.Client to automatically add analytics headers to all requests. + """ + + def __init__( + self, + module_name: str, + analytics_config: tuple[tuple[str, Callable], ...] = ANALYTICS_HEADER_CONFIG, + *args, + **kwargs, + ): + """ + Initialize the AnalyticsHttpClient with analytics headers. + + Parameters + ---------- + module_name : str + The name of the module using this client, added to analytics headers. + analytics_config : tuple[tuple[str, Callable], ...], optional + Configuration for analytics headers, defaults to ANALYTICS_HEADER_CONFIG. + Each entry is a tuple of (header_name, get_value_func) where get_value_func + is a callable that returns the value for the header. + """ + + super().__init__(*args, **kwargs) + self._module_name = module_name + self._analytics_config = analytics_config + + @property + def module_name(self) -> str: + return self._module_name + + def set_module_name(self, module_name: str) -> None: + self._module_name = module_name + + def request(self, method, url, *args, **kwargs): + headers = kwargs.get("headers", {}) + if headers is None: + headers = {} + + new_kwargs = {**kwargs, "headers": self._add_analytics_headers(headers)} + + return super().request(method, url, *args, **new_kwargs) + + def stream(self, method, url, *args, **kwargs): + headers = kwargs.get("headers", {}) + if headers is None: + headers = {} + + new_kwargs = {**kwargs, "headers": self._add_analytics_headers(headers)} + + return super().stream(method, url, *args, **new_kwargs) + + def _add_analytics_headers(self, headers: dict) -> dict: + for header_name, get_value_func in self._analytics_config: + if header_name == ANALYTICS_KEYS.ModuleName: + headers[header_name] = self._module_name + elif get_value_func is not None: + headers[header_name] = get_value_func() + else: + raise ValueError( + f"Invalid analytics config: header_name={header_name}, get_value_func={get_value_func}" + ) + return headers