Add usage analytics utils

requirements.txt

@@ -1,4 +1,5 @@
 numpy>=2.0.0
 pandas>=2.2.3
 scikit-learn>=1.6.1
-typing_extensions>=4.12.2
+typing_extensions>=4.12.2
+httpx~=0.28.1

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()

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",
+]

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
+)

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())

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