feat(helper): actively check version on start (docarray#302)

hanxiao · web-flow · commit 6dadbdff661f · 2022-04-26T11:07:49.000+02:00
diff --git a/docarray/__init__.py b/docarray/__init__.py
@@ -10,3 +10,8 @@
     from rich.traceback import install
 
     install()
+
+if 'NO_VERSION_CHECK' not in os.environ:
+    from .helper import is_latest_version
+
+    is_latest_version()
diff --git a/docarray/helper.py b/docarray/helper.py
@@ -3,9 +3,16 @@
 import pathlib
 import random
 import sys
+import threading
 import uuid
 import warnings
+from distutils.version import LooseVersion
 from typing import Any, Dict, Optional, Sequence, Tuple
+from urllib.request import Request, urlopen
+
+import pkg_resources
+from rich import print
+from rich.panel import Panel
 
 ALLOWED_PROTOCOLS = {'pickle', 'protobuf', 'protobuf-array', 'pickle-array'}
 ALLOWED_COMPRESSIONS = {'lz4', 'bz2', 'lzma', 'zlib', 'gzip'}
@@ -443,3 +450,49 @@ def filter_dict(d: Dict) -> Dict:
     :return: filtered dict
     """
     return dict(filter(lambda item: item[1], d.items()))
+
+
+def _version_check(package: str = None, github_repo: str = None):
+    try:
+        if not package:
+            package = vars(sys.modules[__name__])['__package__']
+        if not github_repo:
+            github_repo = package
+
+        cur_ver = LooseVersion(pkg_resources.get_distribution(package).version)
+        req = Request(
+            f'https://pypi.python.org/pypi/{package}/json',
+            headers={'User-Agent': 'Mozilla/5.0'},
+        )
+        with urlopen(
+            req, timeout=5
+        ) as resp:  # 'with' is important to close the resource after use
+            j = json.load(resp)
+            releases = j.get('releases', {})
+            latest_release_ver = list(
+                sorted(LooseVersion(v) for v in releases.keys() if '.dev' not in v)
+            )[-1]
+            if cur_ver < latest_release_ver:
+                print(
+                    Panel(
+                        f'You are using [b]{package} {cur_ver}[/b], but [bold green]{latest_release_ver}[/] is available. '
+                        f'You may upgrade it via [b]pip install -U {package}[/b]. [link=https://github.com/jina-ai/{github_repo}/blob/main/CHANGELOG.md]Read Changelog here[/link].',
+                        title=':new: New version available!',
+                        width=50,
+                    )
+                )
+    except Exception:
+        # no network, too slow, PyPi is down
+        pass
+
+
+def is_latest_version(package: str = None, github_repo: str = None) -> None:
+    """Check if there is a latest version from Pypi, set env `NO_VERSION_CHECK` to disable it.
+
+    :param package: package name if none auto-detected
+    :param github_repo: repo name that contains CHANGELOG if none then the same as package name
+    """
+
+    threading.Thread(
+        target=_version_check, daemon=True, args=(package, github_repo)
+    ).start()
diff --git a/tests/unit/test_import.py b/tests/unit/test_import.py
@@ -0,0 +1,23 @@
+import multiprocessing
+import threading
+
+
+def test_import_thread_mp():
+    def f1():
+        print(f' f1 start import ')
+        import docarray
+
+        print(f' f1 imported {docarray.__version__}')
+
+    def f2():
+        print(f' f2 start import ')
+        import docarray
+
+        print(f' f2 imported {docarray.__version__}')
+
+    x1 = threading.Thread(target=f1)
+    x2 = multiprocessing.Process(target=f2)
+    x2.start()
+    x1.start()
+    x2.join()
+    x1.join()
