version : modified __getitem__ to return either an item or a Version instance + unit tests

2 files changed, 50 insertions, 7 deletions

diff --git a/lib/spack/spack/test/versions.py b/lib/spack/spack/test/versions.py
index 5f7d57b4e4..c051895ca3 100644
--- a/lib/spack/spack/test/versions.py
+++ b/lib/spack/spack/test/versions.py
@@ -391,9 +391,37 @@ class VersionsTest(unittest.TestCase):
             self.assertEqual(v.underscored, '1_2_3')
 
     def test_repr_and_str(self):
-        a = Version('1.2.3')
-        self.assertEqual(repr(a), 'Version(\'1.2.3\')')
-        b = eval(repr(a))
-        self.assertEqual(a, b)
-        self.assertEqual(str(a), '1.2.3')
-        self.assertEqual(str(a), str(b))
+
+        def check_repr_and_str(vrs):
+            a = Version(vrs)
+            self.assertEqual(repr(a), 'Version(\'' + vrs + '\')')
+            b = eval(repr(a))
+            self.assertEqual(a, b)
+            self.assertEqual(str(a), vrs)
+            self.assertEqual(str(a), str(b))
+
+        check_repr_and_str('1.2.3')
+        check_repr_and_str('R2016a')
+        check_repr_and_str('R2016a.2-3_4')
+
+    def test_get_item(self):
+        a = Version('0.1_2-3')
+        self.assertTrue(isinstance(a[1], int))
+        # Test slicing
+        b = a[0:2]
+        self.assertTrue(isinstance(b, Version))
+        self.assertEqual(b, Version('0.1'))
+        self.assertEqual(repr(b), 'Version(\'0.1\')')
+        self.assertEqual(str(b), '0.1')
+        b = a[0:3]
+        self.assertTrue(isinstance(b, Version))
+        self.assertEqual(b, Version('0.1_2'))
+        self.assertEqual(repr(b), 'Version(\'0.1_2\')')
+        self.assertEqual(str(b), '0.1_2')
+        b = a[1:]
+        self.assertTrue(isinstance(b, Version))
+        self.assertEqual(b, Version('1_2-3'))
+        self.assertEqual(repr(b), 'Version(\'1_2-3\')')
+        self.assertEqual(str(b), '1_2-3')
+        # Raise TypeError on tuples
+        self.assertRaises(TypeError, b.__getitem__, 1, 2)
\ No newline at end of file
diff --git a/lib/spack/spack/version.py b/lib/spack/spack/version.py
index 6f95e8a2f3..7e5878383d 100644
--- a/lib/spack/spack/version.py
+++ b/lib/spack/spack/version.py
@@ -44,6 +44,7 @@ be called on any of the types::
   concrete
 """
 import re
+import numbers
 from bisect import bisect_left
 from functools import wraps
 
@@ -194,7 +195,21 @@ class Version(object):
         return iter(self.version)
 
     def __getitem__(self, idx):
-        return tuple(self.version[idx])
+        cls = type(self)
+        if isinstance(idx, numbers.Integral):
+            return self.version[idx]
+        elif isinstance(idx, slice):
+            # Currently len(self.separators) == len(self.version) - 1
+            extendend_separators = self.separators + ('',)
+            string_arg = []
+            for token, separator in zip(self.version, extendend_separators)[idx]:
+                string_arg.append(str(token))
+                string_arg.append(str(separator))
+            string_arg.pop()  # We don't need the last separator
+            string_arg = ''.join(string_arg)
+            return cls(string_arg)
+        message = '{cls.__name__} indices must be integers'
+        raise TypeError(message.format(cls=cls))
 
     def __repr__(self):
         return 'Version(' + repr(self.string) + ')'