Better checking

apache · Jun 2, 2015 · e89376c · e89376c
1 parent 22703dd
commit e89376c
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/python/pyspark/mllib/__init__.py b/python/pyspark/mllib/__init__.py
@@ -23,14 +23,9 @@
 # MLlib currently needs NumPy 1.4+, so complain if lower
 
 import numpy
-ver = numpy.version.version
 
-fd = ver.find('.')
-ver1 = int(ver[: fd])
-ver = ver[fd + 1:]
-ver2 = int(ver[: ver.find('.')])
-
-if ver1 < 1 or (ver1 == 1 and ver2 <= 4):
+ver = [int(x) for x in numpy.version.version.split(.)[:2]]
+if ver < [1, 4]:
     raise Exception("MLlib requires NumPy 1.4+")
 
 __all__ = ['classification', 'clustering', 'feature', 'fpm', 'linalg', 'random',