Skip to content

Commit 31124e3

Browse files
authored
Implement event scopes (#49)
* Implement event scopes * lint * Add CI-based group id (#48) * Add CI-based group id closes #12 * lint * fix lint
1 parent fc022d1 commit 31124e3

File tree

1 file changed

+63
-2
lines changed

1 file changed

+63
-2
lines changed

src/iterative_telemetry/__init__.py

Lines changed: 63 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
"""Iterative Telemetry."""
2+
import contextlib
3+
import dataclasses
24
import hashlib
35
import json
46
import logging
@@ -7,10 +9,10 @@
79
import subprocess
810
import sys
911
import uuid
10-
from functools import lru_cache
12+
from functools import lru_cache, wraps
1113
from pathlib import Path
1214
from threading import Thread
13-
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
15+
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
1416

1517
import distro
1618
import requests
@@ -28,6 +30,14 @@
2830
DO_NOT_TRACK_VALUE = "do-not-track"
2931

3032

33+
@dataclasses.dataclass
34+
class TelemetryEvent:
35+
interface: str
36+
action: str
37+
error: Optional[str] = None
38+
kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
39+
40+
3141
class IterativeTelemetryLogger:
3242
def __init__(
3343
self,
@@ -47,6 +57,57 @@ def __init__(
4757
if self.debug:
4858
logger.setLevel(logging.DEBUG)
4959
logger.debug("IterativeTelemetryLogger is in debug mode")
60+
self._current_event: Optional[TelemetryEvent] = None
61+
62+
def log_param(self, key: str, value):
63+
if self._current_event:
64+
self._current_event.kwargs[key] = value
65+
66+
@contextlib.contextmanager
67+
def event_scope(
68+
self, interface: str, action: str
69+
) -> Iterator[TelemetryEvent]:
70+
event = TelemetryEvent(interface=interface, action=action)
71+
tmp = self._current_event
72+
self._current_event = event
73+
try:
74+
yield event
75+
finally:
76+
self._current_event = tmp
77+
78+
def log(
79+
self,
80+
interface: str,
81+
action: str = None,
82+
skip: Union[bool, Callable[[TelemetryEvent], bool]] = None,
83+
):
84+
def decorator(func):
85+
@wraps(func)
86+
def inner(*args, **kwargs):
87+
with self.event_scope(
88+
interface, action or func.__name__
89+
) as event:
90+
try:
91+
return func(*args, **kwargs)
92+
except Exception as exc:
93+
event.error = exc.__class__.__name__
94+
raise
95+
finally:
96+
if (
97+
skip is None
98+
or (callable(skip) and not skip(event))
99+
or not skip
100+
):
101+
self.send_event(
102+
event.interface,
103+
event.action,
104+
event.error,
105+
**event.kwargs,
106+
)
107+
108+
return inner
109+
110+
return decorator
50111

51112
def send_cli_call(self, cmd_name: str, error: str = None, **kwargs):
52113
self.send_event("cli", cmd_name, error=error, **kwargs)

0 commit comments

Comments
 (0)