Skip to content

Commit

Permalink
allows disabling the ploomber cloud messsage (#87)
Browse files Browse the repository at this point in the history
* allows disabling cloud message

* fix

* changelog update
  • Loading branch information
edublancas authored Feb 2, 2024
1 parent e8762a6 commit cde2ab3
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 14 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

## 0.2.24dev

* [Feature] Allows disabling the message that promotes Ploomber Cloud when initializing `Telemetry`

## 0.2.23 (2024-01-23)

* [Fix] Fully removes dependency on `click`
Expand Down
32 changes: 21 additions & 11 deletions src/ploomber_core/telemetry/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
telemetry_version - Telemetry version
"""

from copy import copy
from inspect import signature, _empty
import logging
Expand Down Expand Up @@ -326,7 +327,7 @@ def check_cloud():
internal.last_cloud_check = now


def _get_telemetry_info(package_name, version):
def _get_telemetry_info():
"""
The function checks for the local config and uid files, returns the right
values according to the config file (True/False). In addition it checks
Expand All @@ -335,10 +336,6 @@ def _get_telemetry_info(package_name, version):
# Check if telemetry is enabled, if not skip, else check for uid
telemetry_enabled = check_telemetry_enabled()

# Check latest version
check_version(package_name, version)
check_cloud()

if telemetry_enabled:
# Check first time install
is_install = check_first_time_usage()
Expand Down Expand Up @@ -384,7 +381,7 @@ def log_call(self, action=None, payload=False, log_args=False, ignore_args=None)


class Telemetry:
def __init__(self, api_key, package_name, version):
def __init__(self, api_key, package_name, version, *, print_cloud_message=True):
"""
Parameters
Expand All @@ -398,6 +395,9 @@ def __init__(self, api_key, package_name, version):
version : str
Version of the package calling the function
print_cloud_message : bool, default=True
If True, it'll print a message to ask the user to sign up for
Ploomber Cloud
"""
if "_PLOOMBER_TELEMETRY_DEBUG" in os.environ:
warnings.warn(
Expand All @@ -409,16 +409,22 @@ def __init__(self, api_key, package_name, version):
self.api_key = api_key
self.package_name = package_name
self.version = version
self.print_cloud_message = print_cloud_message

@classmethod
def from_package(cls, package_name):
def from_package(cls, package_name, *, print_cloud_message=True):
"""
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)
return cls(
api_key=default_api_key,
package_name=package_name,
version=version,
print_cloud_message=print_cloud_message,
)

def log_api(self, action, client_time=None, total_runtime=None, metadata=None):
"""
Expand All @@ -434,9 +440,13 @@ def log_api(self, action, client_time=None, total_runtime=None, metadata=None):
if client_time is None:
client_time = datetime.datetime.now()

(telemetry_enabled, uid, is_install) = _get_telemetry_info(
self.package_name, self.version
)
(telemetry_enabled, uid, is_install) = _get_telemetry_info()

# Check latest version
check_version(self.package_name, self.version)

if self.print_cloud_message:
check_cloud()

# NOTE: this should not happen anymore
if "NO_UID" in uid:
Expand Down
36 changes: 33 additions & 3 deletions tests/telemetry/test_telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,9 +328,7 @@ def test_full_telemetry_info(monkeypatch, ignore_env_var_and_set_tmp_default_hom
monkeypatch.setattr(telemetry, "DEFAULT_HOME_DIR", str(Path().absolute()))
monkeypatch.setattr(telemetry, "internal", telemetry.Internal())

(stat_enabled, uid, is_install) = telemetry._get_telemetry_info(
"ploomber", "0.14.0"
)
(stat_enabled, uid, is_install) = telemetry._get_telemetry_info()
assert stat_enabled is True
assert isinstance(uid, str)
assert is_install is True
Expand Down Expand Up @@ -1083,3 +1081,35 @@ def test_from_package():
assert _telemetry.api_key
assert _telemetry.version
assert _telemetry.package_name == "ploomber-core"


def test_runs_check_cloud(monkeypatch):
mock = Mock()
monkeypatch.setattr(telemetry, "check_cloud", mock)

my_telemetry = telemetry.Telemetry(
MOCK_API_KEY,
"some-package",
"1.2.2",
print_cloud_message=True,
)

my_telemetry.log_api("some-action")

mock.assert_called_once()


def test_disable_check_cloud(monkeypatch):
mock = Mock()
monkeypatch.setattr(telemetry, "check_cloud", mock)

my_telemetry = telemetry.Telemetry(
MOCK_API_KEY,
"some-package",
"1.2.2",
print_cloud_message=False,
)

my_telemetry.log_api("some-action")

mock.assert_not_called()

0 comments on commit cde2ab3

Please sign in to comment.