Update version detection script for CI

Signed-off-by: Joffrey F <joffrey@docker.com>

1 files changed, 24 insertions, 15 deletions

diff --git a/scripts/versions.py b/scripts/versions.py
index 77aaf4f..7212195 100644
--- a/scripts/versions.py
+++ b/scripts/versions.py
@@ -11,23 +11,24 @@ categories = [
     'test'
 ]
 
+STAGES = ['tp', 'beta', 'rc']
 
-class Version(namedtuple('_Version', 'major minor patch rc edition')):
+
+class Version(namedtuple('_Version', 'major minor patch stage edition')):
 
     @classmethod
     def parse(cls, version):
         edition = None
         version = version.lstrip('v')
-        version, _, rc = version.partition('-')
-        if rc:
-            if 'rc' not in rc:
-                edition = rc
-                rc = None
-            elif '-' in rc:
-                edition, rc = rc.split('-')
-
+        version, _, stage = version.partition('-')
+        if stage:
+            if not any(marker in stage for marker in STAGES):
+                edition = stage
+                stage = None
+            elif '-' in stage:
+                edition, stage = stage.split('-')
         major, minor, patch = version.split('.', 3)
-        return cls(major, minor, patch, rc, edition)
+        return cls(major, minor, patch, stage, edition)
 
     @property
     def major_minor(self):
@@ -38,14 +39,22 @@ class Version(namedtuple('_Version', 'major minor patch rc edition')):
         """Return a representation that allows this object to be sorted
         correctly with the default comparator.
         """
-        # rc releases should appear before official releases
-        rc = (0, self.rc) if self.rc else (1, )
-        return (int(self.major), int(self.minor), int(self.patch)) + rc
+        # non-GA releases should appear before GA releases
+        # Order: tp -> beta -> rc -> GA
+        if self.stage:
+            for st in STAGES:
+                if st in self.stage:
+                    stage = (STAGES.index(st), self.stage)
+                    break
+        else:
+            stage = (len(STAGES),)
+
+        return (int(self.major), int(self.minor), int(self.patch)) + stage
 
     def __str__(self):
-        rc = '-{}'.format(self.rc) if self.rc else ''
+        stage = '-{}'.format(self.stage) if self.stage else ''
         edition = '-{}'.format(self.edition) if self.edition else ''
-        return '.'.join(map(str, self[:3])) + edition + rc
+        return '.'.join(map(str, self[:3])) + edition + stage
 
 
 def main():