11"""Iterative Telemetry."""
2+ import contextlib
3+ import dataclasses
24import hashlib
35import json
46import logging
79import subprocess
810import sys
911import uuid
10- from functools import lru_cache
12+ from functools import lru_cache , wraps
1113from pathlib import Path
1214from 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
1517import distro
1618import requests
2830DO_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+
3141class 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