Version update - New branch

doc/community/user-stats.rst

@@ -1,5 +1,5 @@
 User Statistics
-==========
+===============
 
 As an open source project, we collect anonymous usage statistics to prioritize and find product gaps.
 This is optional and may be turned off by changing the configuration file:
@@ -12,3 +12,11 @@ The data we collect is limited to:
 - A generated UUID, randomized when the initial install takes place, no personal or any identifiable information.
 - Environment variables: The OS architecture Ploomber is used in (Python version etc.)
 - Information about the different product phases: installation, API calls and errors.
+
+Version updates
+---------------
+If there's an outdated version, ploomber will alert it through the console every second day in a non-invasive way.
+You can stop this checks for instance if you're running in production and you've locked versions.
+The check can be turned off by changing the configuration file:
+ inside ~/.ploomber/stats/config.yaml
+ Change version_check_enabled to False.

src/ploomber/telemetry/telemetry.py

@@ -31,8 +31,10 @@
 
 import datetime
 import http.client as httplib
+import json
 import warnings
 
+import click
 import posthog
 import yaml
 import os
@@ -254,21 +256,14 @@ def check_uid():
     uid_path = Path(check_dir_exist(CONF_DIR), 'uid.yaml')
     if not uid_path.exists():  # Create - doesn't exist
         uid = str(uuid.uuid4())
-        try:  # Create for future runs
-            with uid_path.open("w") as file:
-                yaml.dump({"uid": uid}, file)
+        res = write_conf_file(uid_path, {"uid": uid}, error=True)
+        if res:
+            return f"NO_UID {res}"
+        else:
             return uid
-        except Exception as e:
-            warnings.warn(f"ERROR: Can't write UID file: {e}")
-            return f"NO_UID {e}"
     else:  # read and return uid
-        try:
-            with uid_path.open("r") as file:
-                uid_dict = yaml.safe_load(file)
-            return uid_dict['uid']
-        except Exception as e:
-            warnings.warn(f"Error: Can't read UID file: {e}")
-            return f"NO_UID {e}"
+        conf = read_conf_file(uid_path)
+        return conf.get('uid', "NO_UID")
 
 
 def check_stats_enabled():
@@ -284,21 +279,11 @@ def check_stats_enabled():
     # Check if local config exists
     config_path = Path(check_dir_exist(CONF_DIR), 'config.yaml')
     if not config_path.exists():
-        try:  # Create for future runs
-            with config_path.open("w") as file:
-                yaml.dump({"stats_enabled": True}, file)
-            return True
-        except Exception as e:
-            warnings.warn(f"ERROR: Can't write to config file: {e}")
-            return True
+        write_conf_file(config_path, {"stats_enabled": True})
+        return True
     else:  # read and return config
-        try:
-            with config_path.open("r") as file:
-                conf = yaml.safe_load(file)
-            return conf['stats_enabled']
-        except Exception as e:
-            warnings.warn(f"Error: Can't read config file {e}")
-            return True
+        conf = read_conf_file(config_path)
+        return conf.get('stats_enabled', True)
 
 
 def check_first_time_usage():
@@ -311,6 +296,86 @@ def check_first_time_usage():
     return not uid_path.exists() and config_path.exists()
 
 
+def get_latest_version():
+    """
+    The function checks for the latest available ploomber version
+    uid file doesn't exist.
+    """
+    conn = httplib.HTTPSConnection('pypi.org', timeout=1)
+    try:
+        conn.request("GET", "/pypi/ploomber/json")
+        content = conn.getresponse().read()
+        data = json.loads(content)
+        latest = data['info']['version']
+        return latest
+    except Exception:
+        return __version__
+    finally:
+        conn.close()
+
+
+def read_conf_file(conf_path):
+    try:
+        with conf_path.open("r") as file:
+            conf = yaml.safe_load(file)
+    except Exception as e:
+        warnings.warn(f"Error: Can't read config file {e}")
+    return conf
+
+
+def write_conf_file(conf_path, to_write, error=None):
+    try:  # Create for future runs
+        with conf_path.open("w") as file:
+            yaml.dump(to_write, file)
+    except Exception as e:
+        warnings.warn(f"ERROR: Can't write to config file: {e}")
+        if error:
+            return e
+
+
+def check_version():
+    """
+    The function checks if the user runs the latest version
+    This check will be skipped if the version_check_enabled is set to False
+    If it's not the latest, notifies the user and saves the metadata to conf
+    Alerting every 2 days on stale versions
+    """
+    # Read conf file
+    config_path = Path(check_dir_exist(CONF_DIR), 'config.yaml')
+    conf = read_conf_file(config_path)
+    # Check if the flag was disabled
+    if conf and 'version_check_enabled' in conf.keys() \
+            and not conf['version_check_enabled']:
+        return
+
+    # If latest version, do nothing
+    latest = get_latest_version()
+    version = __version__
+    print(version)
+    if __version__ == latest:
+        return
+
+    today = datetime.datetime.now()
+
+    # First time the version is checked
+    if 'version_check_enabled' in conf.keys():
+
+        # Check if we already notified in the last 2 days
+        last_message = conf['version_check_enabled']
+        diff = (today - last_message).days
+        if diff < 2:
+            return
+
+    click.secho(
+        "Please update the ploomber version, run:\n"
+        "pip install ploomber --upgrade\n",
+        fg='yellow')
+
+    # Update conf
+    conf['version_check_enabled'] = today
+    write_conf_file(config_path, conf)
+
+
 def _get_telemetry_info():
     """
     The function checks for the local config and uid files, returns the right
@@ -321,6 +386,9 @@ def _get_telemetry_info():
     telemetry_enabled = check_stats_enabled()
 
     if telemetry_enabled:
+        # Check latest version
+        check_version()
+
         # Check first time install
         is_install = check_first_time_usage()
 
@@ -423,9 +491,7 @@ def log_api(action,
 def log_exception(action):
     """Runs a function and logs exceptions, if any
     """
-
     def _log_exceptions(func):
-
         @wraps(func)
         def wrapper(*args, **kwargs):
             try:

tests/telemetry/test_telemetry.py

@@ -1,11 +1,14 @@
+import datetime
 import pathlib
+import click
 import sys
 from unittest.mock import Mock
 from pathlib import Path
 
-import click
 import pytest
+import yaml
 
+import ploomber
 from ploomber.telemetry import telemetry
 from ploomber.telemetry.validate_inputs import str_param, opt_str_param
 from ploomber.cli import plot, install, build, interact, task, report, status
@@ -144,7 +147,6 @@ def test_pip_env(monkeypatch, inside_pip_env):
 # Ref: https://stackoverflow.com/questions/43878953/how-does-one-detect-if-
 # one-is-running-within-a-docker-container-within-python
 def test_docker_env(monkeypatch):
-
     def mock(input_path):
         return 'dockerenv' in str(input_path)
 
@@ -392,7 +394,6 @@ def test_parse_dag_products(monkeypatch):
 
 
 def test_parse_dag(monkeypatch, tmp_directory):
-
     def fn1(product):
         pass
 
@@ -464,6 +465,89 @@ def test_validate_entries(monkeypatch):
     assert res == (event_id, uid, action, client_time, elapsed_time)
 
 
+def test_get_latest_version(monkeypatch):
+    is_latest = telemetry.get_latest_version()
+    assert isinstance(is_latest, str)
+    version_index = [i for i, ltr in enumerate(is_latest) if ltr == '.']
+    assert len(version_index) >= 1
+
+    mock_httplib = Mock()
+    mock_httplib.HTTPSConnection().request.side_effect = Exception
+    monkeypatch.setattr(telemetry, 'httplib', mock_httplib)
+    is_latest = telemetry.get_latest_version()
+    assert is_latest == ploomber.__version__
+
+    # Mock version and the conf, check it produces the same version
+
+
+def write_to_conf_file(tmp_directory, monkeypatch, last_check):
+    stats = Path('stats')
+    stats.mkdir()
+    path = stats / 'config.yaml'
+    monkeypatch.setattr(telemetry, 'DEFAULT_HOME_DIR', '.')
+    (path).write_text("stats_enabled: True\n"
+                      f"version_check_enabled: {last_check}\n")
+
+
+def test_version_skips_when_updated(tmp_directory, capsys, monkeypatch):
+    # Path conf file
+    monkeypatch.setattr(telemetry, '__version__', '0.14.8')
+    mock_version = Mock()
+    mock_version.return_value = '0.14.8'
+    monkeypatch.setattr(telemetry, 'get_latest_version', mock_version)
+
+    write_to_conf_file(
+        tmp_directory=tmp_directory,
+        monkeypatch=monkeypatch,
+        last_check='2022-01-20 10:51:41.082376')  # version='0.14.8',
+
+    # Test no warning when same version encountered
+    telemetry.check_version()
+    captured = capsys.readouterr()
+    assert "ploomber version" not in captured.out
+
+
+def test_user_output_on_different_versions(tmp_directory, capsys, monkeypatch):
+    mock_version = Mock()
+    monkeypatch.setattr(telemetry, 'get_latest_version', mock_version)
+    write_to_conf_file(tmp_directory=tmp_directory,
+                       monkeypatch=monkeypatch,
+                       last_check='2022-01-20 10:51:41.082376')
+    mock_version.return_value = '0.14.0'
+
+    # Check now that the date is different there is an upgrade warning
+    telemetry.check_version()
+    captured = capsys.readouterr()
+    assert "ploomber version" in captured.out
+
+
+def test_no_output_latest_version(tmp_directory, capsys, monkeypatch):
+    # The file's date is today now, no error should be raised
+    write_to_conf_file(tmp_directory=tmp_directory,
+                       monkeypatch=monkeypatch,
+                       last_check=datetime.datetime.now())
+    telemetry.check_version()
+    captured = capsys.readouterr()
+    assert "ploomber version" not in captured.out
+
+
+def test_output_on_date_diff(tmp_directory, capsys, monkeypatch):
+    # Warning should be caught since the date and version are off
+    write_to_conf_file(tmp_directory=tmp_directory,
+                       monkeypatch=monkeypatch,
+                       last_check='2022-01-20 10:51:41.082376')
+    path = Path('stats') / 'config.yaml'
+    telemetry.check_version()
+    captured = capsys.readouterr()
+    assert "ploomber version" in captured.out
+
+    # Check the conf file was updated
+    with path.open("r") as file:
+        conf = yaml.safe_load(file)
+    diff = (datetime.datetime.now() - conf['version_check_enabled']).days
+    assert diff == 0
+
+
 def test_python_major_version():
     version = telemetry.python_version()
     major = version.split(".")[0]