faster telemetry + simpler initialization (#82)

* more efficient system info

* adds classmethod to initialize telemetry with default settings

* lint

* changelog and docs

* update doc

* ci

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.2.20dev
 
+* [Feature] Obtain system info once for faster initialization
+* [Feature] Add `Telemetry.from_package` for simpler initialization
+
+
 ## 0.2.19 (2023-12-12)
 
 * [Fix] `@modify_exceptions` decorator now compatible with `ClickException` ([#81](https://github.com/ploomber/core/issues/81))

doc/telemetry.md

@@ -27,16 +27,17 @@ from ploomber_core.telemetry import Telemetry
 
 Initialize it with the API key, name of the package, and version. Here's an example:
 
+```{versionadded} 0.2.20
+`Telemetry.from_package`
+```
+
 ```{code-cell} ipython3
-from ploomber_core import __version__
-
-telemetry = Telemetry(
-    api_key="SOMEKEY",
-    # change for the actual name
-    package_name="ploomber-core",
-    # change for the actual version
-    version=__version__,
-)
+telemetry = Telemetry.from_package(package_name="ploomber-core")
+```
+
+```{note}
+`Telemetry.from_package` is the simplest way to initialize the telemetry object, you might use the constructor directly if you want to customize
+the configuration.
 ```
 
 To log call functions:

setup.py

@@ -17,6 +17,7 @@
     "click",
     "pyyaml",
     "posthog",
+    'importlib-metadata;python_version<"3.8"',
 ]
 
 DEV = [

src/ploomber_core/telemetry/system_info.py

@@ -4,6 +4,12 @@
 import sys
 
 
+try:
+    import importlib.metadata as importlib_metadata
+except ModuleNotFoundError:
+    import importlib_metadata
+
+
 def python_version():
     py_version = sys.version_info
     return f"{py_version.major}.{py_version.minor}.{py_version.micro}"
@@ -106,3 +112,31 @@ def is_airflow():
 def is_argo():
     """Returns: True for Argo env"""
     return "ARGO_AGENT_TASK_WORKERS" in os.environ or "ARGO_KUBELET_PORT" in os.environ
+
+
+def safe_call(function):
+    try:
+        return function()
+    except Exception:
+        return None
+
+
+def get_package_version(package_name):
+    try:
+        return importlib_metadata.version(package_name)
+    except importlib_metadata.PackageNotFoundError:
+        return None
+
+
+def get_system_info():
+    return {
+        "os": safe_call(get_os),
+        "python_version": safe_call(python_version),
+        "env": safe_call(get_env),
+        "is_docker": safe_call(is_docker),
+        "is_colab": safe_call(is_colab),
+        "is_paperspace": safe_call(is_paperspace),
+        "is_slurm": safe_call(is_slurm),
+        "is_airflow": safe_call(is_airflow),
+        "is_argo": safe_call(is_argo),
+    }

src/ploomber_core/telemetry/telemetry.py

@@ -46,17 +46,7 @@
 
 from ploomber_core.telemetry import validate_inputs
 from ploomber_core.config import Config
-from ploomber_core.telemetry.system_info import (
-    get_env,
-    get_os,
-    is_airflow,
-    is_argo,
-    is_colab,
-    is_docker,
-    is_paperspace,
-    is_slurm,
-    python_version,
-)
+from ploomber_core.telemetry.system_info import get_system_info, get_package_version
 
 TELEMETRY_VERSION = "0.5"
 DEFAULT_HOME_DIR = str(Path.home() / ".ploomber")
@@ -69,6 +59,9 @@
 logging.getLogger("posthog").disabled = True
 
 
+SYSTEM_INFO = get_system_info()
+
+
 class UserSettings(Config):
     """User-customizable settings"""
 
@@ -420,6 +413,16 @@ def __init__(self, api_key, package_name, version):
         self.package_name = package_name
         self.version = version
 
+    @classmethod
+    def from_package(cls, package_name):
+        """
+        Initialize a Telemetry client with the default configuration for
+        a package with the given name
+        """
+        default_api_key = "phc_P9SpSeypyPwxrMdFn2edOOEooQioF2axppyEeDwtMSP"
+        version = get_package_version(package_name)
+        return cls(api_key=default_api_key, package_name=package_name, version=version)
+
     def log_api(self, action, client_time=None, total_runtime=None, metadata=None):
         """
         This function logs through an API call, assigns parameters
@@ -443,35 +446,38 @@ def log_api(self, action, client_time=None, total_runtime=None, metadata=None):
             metadata["uid_issue"] = uid
             uid = None
 
-        py_version = python_version()
-        docker_container = is_docker()
         cloud = is_cloud_user()
         email = email_registered()
-        colab = is_colab()
+        colab = SYSTEM_INFO["is_colab"]
+
         if colab:
             metadata["colab"] = colab
 
-        paperspace = is_paperspace()
+        paperspace = SYSTEM_INFO["is_paperspace"]
+
         if paperspace:
             metadata["paperspace"] = paperspace
 
-        slurm = is_slurm()
+        slurm = SYSTEM_INFO["is_slurm"]
+
         if slurm:
             metadata["slurm"] = slurm
 
-        airflow = is_airflow()
+        airflow = SYSTEM_INFO["is_airflow"]
+
         if airflow:
             metadata["airflow"] = airflow
 
-        argo = is_argo()
+        argo = SYSTEM_INFO["is_argo"]
+
         if argo:
             metadata["argo"] = argo
 
         if "dag" in metadata:
             metadata["dag"] = parse_dag(metadata["dag"])
 
-        os = get_os()
-        environment = get_env()
+        os = SYSTEM_INFO["os"]
+        environment = SYSTEM_INFO["env"]
 
         if telemetry_enabled:
             (event_id, uid, action, client_time, elapsed_time) = validate_entries(
@@ -483,10 +489,10 @@ def log_api(self, action, client_time=None, total_runtime=None, metadata=None):
                 "action": action,
                 "client_time": str(client_time),
                 "total_runtime": total_runtime,
-                "python_version": py_version,
+                "python_version": SYSTEM_INFO["python_version"],
                 "version": self.version,
                 "package_name": self.package_name,
-                "docker_container": docker_container,
+                "docker_container": SYSTEM_INFO["is_docker"],
                 "cloud": cloud,
                 "email": email,
                 "os": os,

tests/telemetry/test_system_info.py

@@ -0,0 +1,15 @@
+from ploomber_core.telemetry.system_info import get_system_info
+
+
+def test_get_system_info():
+    assert set(get_system_info()) == {
+        "env",
+        "is_airflow",
+        "is_argo",
+        "is_colab",
+        "is_docker",
+        "is_paperspace",
+        "is_slurm",
+        "os",
+        "python_version",
+    }

tests/telemetry/test_telemetry.py

@@ -16,6 +16,7 @@
 from ploomber_core.telemetry import system_info
 from ploomber_core.telemetry.validate_inputs import str_param, opt_str_param
 
+
 from ploomber_core.exceptions import BaseException
 
 MOCK_API_KEY = "phc_P1dsjk20bijsabdaib2eu"
@@ -1074,3 +1075,11 @@ def test_check_cloud(tmp_directory, monkeypatch, capsys, last_cloud_check):
 
     assert expected in captured.out
     assert config["last_cloud_check"] == now
+
+
+def test_from_package():
+    _telemetry = telemetry.Telemetry.from_package("ploomber-core")
+
+    assert _telemetry.api_key
+    assert _telemetry.version
+    assert _telemetry.package_name == "ploomber-core"