Get the version from PKG-INFO if it exists

1 files changed, 19 insertions, 0 deletions

diff --git a/setup.py b/setup.py
index f844d52..a99d044 100755
--- a/setup.py
+++ b/setup.py
@@ -2,6 +2,7 @@
 """Distutils installer for testtools."""
 
 from distutils.core import setup
+import email
 import os
 
 import testtools
@@ -13,12 +14,30 @@ def get_revno():
     return t.branch.revno()
 
 
+def get_version_from_pkg_info():
+    """Get the version from PKG-INFO file if we can."""
+    pkg_info_path = os.path.join(os.path.dirname(__file__), 'PKG-INFO')
+    try:
+        pkg_info_file = open(pkg_info_path, 'r')
+    except (IOError, OSError):
+        return None
+    try:
+        pkg_info = email.message_from_file(pkg_info_file)
+    except email.MessageError:
+        return None
+    return pkg_info.get('Version', None)
+
+
 def get_version():
+    """Return the version of testtools that we are building."""
     version = '.'.join(
         str(component) for component in testtools.__version__[0:3])
     phase = testtools.__version__[3]
     if phase == 'final':
         return version
+    pkg_info_version = get_version_from_pkg_info()
+    if pkg_info_version:
+        return pkg_info_version
     revno = get_revno()
     if phase == 'alpha':
         # No idea what the next version will be