Use distutils version comparison.

I also reverted back to use if statements, which felt simpler in this (2
option) case.

Author: drvinceknight

axelrod/plot.py

@@ -3,6 +3,7 @@
 import tqdm
 import warnings
 
+from distutils.version import LooseVersion
 from typing import List, Union
 
 matplotlib_installed = True
@@ -25,18 +26,11 @@
 dataType = List[List[Union[int, float]]]
 
 
-def default_cmap(version: str = None) -> str:
+def default_cmap(version: str = "2.0") -> str:
     """Sets a default matplotlib colormap based on the version."""
-    version_gt_1_5 = False
-    if version:
-        s = version.split('.')
-        version_gt_1_5 = (int(s[0]) == 1 and int(s[1]) >= 5) or (int(s[0]) > 1)
-
-    cmaps = {
-        True: 'viridis',
-        False: 'YlGnBu'
-    }
-    return cmaps[version_gt_1_5]
+    if LooseVersion(version) >= "1.5":
+        return 'viridis'
+    return 'YlGnBu'
 
 
 class Plot(object):

axelrod/tests/unit/test_plot.py

@@ -74,15 +74,24 @@ def setUpClass(cls):
             ['Defector', 'Tit For Tat', 'Alternator'])
 
     def test_default_cmap(self):
-        cmap = axelrod.plot.default_cmap()
+        cmap = axelrod.plot.default_cmap('0.0')
         self.assertEqual(cmap, 'YlGnBu')
 
-        cmap = axelrod.plot.default_cmap('0.0')
+        cmap = axelrod.plot.default_cmap('1.3alpha')
         self.assertEqual(cmap, 'YlGnBu')
 
         cmap = axelrod.plot.default_cmap('1.4')
         self.assertEqual(cmap, 'YlGnBu')
 
+        cmap = axelrod.plot.default_cmap()
+        self.assertEqual(cmap, 'viridis')
+
+        cmap = axelrod.plot.default_cmap('1.5')
+        self.assertEqual(cmap, 'viridis')
+
+        cmap = axelrod.plot.default_cmap('1.5beta')
+        self.assertEqual(cmap, 'viridis')
+
         cmap = axelrod.plot.default_cmap('1.7')
         self.assertEqual(cmap, 'viridis')