Allow for more complex version comparison

dbt/version.py

@@ -1,5 +1,4 @@
-import argparse
-import os
+from string import ascii_lowercase
 import re
 
 try:
@@ -14,59 +13,108 @@
     'master/.bumpversion.cfg'
 
 
-def __parse_version(contents):
-    matches = re.search(r"current_version = ([\.0-9]+)", contents)
-    if matches is None or len(matches.groups()) != 1:
-        return "unknown"
-    else:
-        version = matches.groups()[0]
-        return version
+class Version(object):
 
+    def __init__(self, name, is_latest=True):
+        self.name = name
+        self.parts = self.get_parts()
+        self.numeric_parts = self.get_numeric_parts()
+        self.is_latest = is_latest
 
-def get_version():
-    return __version__
+    def get_parts(self):
+        return self.name.split('.')
 
+    def get_numeric_parts(self):
+        numeric_name = ''.join(
+            [ch for ch in self.name if ch not in ascii_lowercase])
+        return numeric_name.split('.')
 
-def get_latest_version():
+    def __str__(self):
+        if self.is_latest:
+            return "Current version: {}\n".format(self.name)
+        return "Installed version: {}\n".format(self.name)
+
+    def __eq__(self, other):
+        if not isinstance(other, Version):
+            return NotImplemented
+        return self.name == other.name
+
+    def __gt__(self, other):
+        if not isinstance(other, Version):
+            return NotImplemented
+
+        if len(self.parts) > len(other.parts):
+            return True
+
+        zipped_parts = zip(self.numeric_parts, other.numeric_parts)
+        for self_part, other_part in zipped_parts:
+            if other_part > self_part:
+                return False
+            if self_part > other_part:
+                return True
+        return False
+
+    def __lt__(self, other):
+        if not isinstance(other, Version):
+            return NotImplemented
+
+        if len(self.parts) < len(other.parts):
+            return True
+
+        zipped_parts = zip(self.numeric_parts, other.numeric_parts)
+        for self_part, other_part in zipped_parts:
+            if other_part > self_part:
+                return True
+            if self_part > other_part:
+                return False
+        return False
+
+
+def get_version_string_from_text(contents):
+    matches = re.search(r"current_version = ([\.0-9a-z]+)", contents)
+    if matches is None or len(matches.groups()) != 1:
+        return ""
+    version = matches.groups()[0]
+    return version
+
+
+def get_remote_version_file_contents(url=REMOTE_VERSION_FILE):
     try:
-        f = urlopen(REMOTE_VERSION_FILE)
+        f = urlopen(url)
         contents = f.read()
     except:
         contents = ''
     if hasattr(contents, 'decode'):
         contents = contents.decode('utf-8')
-    return __parse_version(contents)
-
+    return contents
 
-def not_latest():
-    return """Your version of dbt is out of date! You can find instructions
-    for upgrading here:
 
-    https://docs.getdbt.com/docs/installation
-    """
+def get_latest_version():
+    contents = get_remote_version_file_contents()
+    version_string = get_version_string_from_text(contents)
+    return Version(version_string)
 
 
-def get_version_string():
-    return "installed version: {}\n   latest version: {}".format(
-        installed, latest
-    )
+def get_installed_version():
+    return Version(__version__, is_latest=False)
 
 
 def get_version_information():
-    basic = get_version_string()
-
-    if is_latest():
-        basic += '\nUp to date!'
+    installed = get_installed_version()
+    latest = get_latest_version()
+    if installed == latest:
+        return "{}{}Up to date!".format(installed, latest)
+
+    elif installed > latest:
+        return "{}{}Your version is ahead!".format(
+                installed, latest)
     else:
-        basic += '\n{}'.format(not_latest())
-
-    return basic
-
-
-def is_latest():
-    return installed == latest
+        return "{}{}Your version of dbt is out of date!\n" \
+            "\tYou can find instructions for upgrading here:\n" \
+            "\thttps://docs.getdbt.com/docs/installation" \
+            .format(installed, latest)
 
 
 __version__ = '0.9.0'
-installed = get_version()
+installed = get_installed_version()
 latest = get_latest_version()

test/unit/test_version.py

@@ -0,0 +1,92 @@
+from mock import patch, MagicMock
+import unittest
+
+import dbt.version
+
+
+class VersionTest(unittest.TestCase):
+
+    @patch("dbt.version.__version__", "1.1.1")
+    def test_versions_equal(self):
+
+        dbt.version.get_remote_version_file_contents = MagicMock(
+            return_value="""
+                [bumpversion]
+                current_version = 1.1.1
+                commit = True
+                tag = True
+
+                [bumpversion:file:setup.py]
+
+                [bumpversion:file:dbt/version.py]
+            """)
+
+        latest_version = dbt.version.get_latest_version()
+        installed_version = dbt.version.get_installed_version()
+        version_information = dbt.version.get_version_information()
+
+        expected_version_information = "Installed version: 1.1.1\n" \
+            "Current version: 1.1.1\n" \
+            "Up to date!"
+
+        assert isinstance(latest_version, dbt.version.Version)
+        assert isinstance(installed_version, dbt.version.Version)
+        self.assertTrue(latest_version.is_latest)
+        self.assertFalse(installed_version.is_latest)
+        self.assertEqual(latest_version, installed_version)
+        self.assertMultiLineEqual(version_information,
+                                  expected_version_information)
+
+    @patch("dbt.version.__version__", "1.12.1")
+    def test_installed_version_greater(self):
+        dbt.version.get_remote_version_file_contents = MagicMock(
+            return_value="""
+                [bumpversion]
+                current_version = 1.1.12
+                commit = True
+                tag = True
+
+                [bumpversion:file:setup.py]
+
+                [bumpversion:file:dbt/version.py]
+            """)
+
+        latest_version = dbt.version.get_latest_version()
+        installed_version = dbt.version.get_installed_version()
+        version_information = dbt.version.get_version_information()
+
+        expected_version_information = "Installed version: 1.12.1\n" \
+            "Current version: 1.1.12\n" \
+            "Your version is ahead!"
+
+        assert installed_version > latest_version
+        self.assertMultiLineEqual(version_information,
+                                  expected_version_information)
+
+    @patch("dbt.version.__version__", "1.10.1a")
+    def test_installed_version_lower(self):
+        dbt.version.get_remote_version_file_contents = MagicMock(
+            return_value="""
+                [bumpversion]
+                current_version = 2.0.113a
+                commit = True
+                tag = True
+
+                [bumpversion:file:setup.py]
+
+                [bumpversion:file:dbt/version.py]
+            """)
+
+        latest_version = dbt.version.get_latest_version()
+        installed_version = dbt.version.get_installed_version()
+        version_information = dbt.version.get_version_information()
+
+        expected_version_information = "Installed version: 1.10.1a\n" \
+            "Current version: 2.0.113a\n" \
+            "Your version of dbt is out of date!\n" \
+            "\tYou can find instructions for upgrading here:\n" \
+            "\thttps://docs.getdbt.com/docs/installation"
+
+        assert installed_version < latest_version
+        self.assertMultiLineEqual(version_information,
+                                  expected_version_information)