Refactor click instrumentation.

cli_telemetry/instrumentation/__init__.py

@@ -0,0 +1,21 @@
+"""
+Automatic instrumentation registry for supported libraries.
+"""
+import os
+
+from .click import auto_instrument_click
+
+def init_auto_instrumentation() -> None:
+    """
+    Initialize automatic instrumentation for all supported libraries.
+    Future instrumentation modules should be added here.
+    """
+    # Click instrumentation
+    if "CLI_TELEMETRY_DISABLE_CLICK_INSTRUMENTATION" in os.environ:
+        return
+    try:
+        auto_instrument_click()
+    except Exception:
+        # Safely ignore instrumentation errors
+        pass
+

cli_telemetry/instrumentation/click.py

@@ -0,0 +1,40 @@
+"""
+Instrumentation for Click commands to auto-wrap their invocation in telemetry spans.
+"""
+
+from ..telemetry import Span, add_tag
+
+def auto_instrument_click():
+    """Monkeypatch click.Command and click.Group to auto-wrap in telemetry spans."""
+    import os
+
+    if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH") == "1":
+        return  # User has opted out of monkeypatching
+
+    try:
+        import click
+
+        def telemetry_wrapper(original_invoke):
+            def invoke_with_span(self, ctx):
+                # Assumes init_telemetry() already called by app
+                with Span(self.name, attributes={"cli.command": ctx.command_path}) as span:
+                    try:
+                        for param in self.params:
+                            if param.name in ctx.params:
+                                add_tag(f"args.{param.name}", ctx.params[param.name])
+                    except Exception:
+                        pass
+                    return original_invoke(self, ctx)
+            return invoke_with_span
+
+        # Avoid double-patching
+        if not getattr(click.Command, "_telemetry_patched", False):
+            click.Command.invoke = telemetry_wrapper(click.Command.invoke)
+            click.Command._telemetry_patched = True
+
+        if not getattr(click.Group, "_telemetry_patched", False):
+            click.Group.invoke = telemetry_wrapper(click.Group.invoke)
+            click.Group._telemetry_patched = True
+
+    except ImportError:
+        pass  # No click? No telemetry sadness.

cli_telemetry/telemetry.py

@@ -69,6 +69,7 @@ def _init_db_file(db_file: str) -> None:
 
 
 def init_telemetry(service_name: str, db_path: Optional[str] = None, user_id_file: Optional[str] = None) -> None:
+    from .instrumentation import init_auto_instrumentation
     """
     Initialize trace ID, user‐ID file, and SQLite DB.
     If db_path/user_id_file are provided, uses those; otherwise defaults to:
@@ -81,9 +82,7 @@ def init_telemetry(service_name: str, db_path: Optional[str] = None, user_id_fil
         if _initialized:
             return
 
-        # monkeypatch click if allowed
-        if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH"):
-            _monkeypatch_click()
+        init_auto_instrumentation()
 
         # determine base path under XDG_DATA_HOME
         xdg = os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
@@ -272,37 +271,3 @@ def end_session() -> None:
             pass
 
 
-def _monkeypatch_click():
-    """Monkeypatch click.Command and click.Group to auto-wrap in telemetry spans."""
-    import os
-
-    if os.environ.get("CLI_TELEMETRY_DISABLE_CLICK_PATCH") == "1":
-        return  # User has opted out of monkeypatching
-
-    try:
-        import click
-
-        def telemetry_wrapper(original_invoke):
-            def invoke_with_span(self, ctx):
-                # Assumes init_telemetry() already called by app
-                with Span(self.name, attributes={"cli.command": ctx.command_path}) as span:
-                    try:
-                        for param in self.params:
-                            if param.name in ctx.params:
-                                add_tag(f"args.{param.name}", ctx.params[param.name])
-                    except Exception:
-                        pass
-                    return original_invoke(self, ctx)
-            return invoke_with_span
-
-        # Avoid double-patching
-        if not getattr(click.Command, "_telemetry_patched", False):
-            click.Command.invoke = telemetry_wrapper(click.Command.invoke)
-            click.Command._telemetry_patched = True
-
-        if not getattr(click.Group, "_telemetry_patched", False):
-            click.Group.invoke = telemetry_wrapper(click.Group.invoke)
-            click.Group._telemetry_patched = True
-
-    except ImportError:
-        pass  # No click? No telemetry sadness.