Clean unittest2 imports up added Python 2.6 before

HyukjinKwon · HyukjinKwon · commit ed834f5062b1 · 2018-11-18T16:17:48.000+08:00
diff --git a/python/pyspark/ml/tests/test_algorithms.py b/python/pyspark/ml/tests/test_algorithms.py
@@ -16,17 +16,8 @@
 #
 
 from shutil import rmtree
-import sys
 import tempfile
-
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 import numpy as np
 
diff --git a/python/pyspark/ml/tests/test_base.py b/python/pyspark/ml/tests/test_base.py
@@ -15,15 +15,7 @@
 # limitations under the License.
 #
 
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 from pyspark.sql.types import DoubleType, IntegerType
 from pyspark.testing.mlutils import MockDataset, MockEstimator, MockUnaryTransformer, \
diff --git a/python/pyspark/ml/tests/test_evaluation.py b/python/pyspark/ml/tests/test_evaluation.py
@@ -15,15 +15,7 @@
 # limitations under the License.
 #
 
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 import numpy as np
 
diff --git a/python/pyspark/ml/tests/test_feature.py b/python/pyspark/ml/tests/test_feature.py
@@ -17,14 +17,7 @@
 #
 
 import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 if sys.version > '3':
     basestring = str
diff --git a/python/pyspark/ml/tests/test_image.py b/python/pyspark/ml/tests/test_image.py
@@ -14,15 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 import py4j
 
diff --git a/python/pyspark/ml/tests/test_linalg.py b/python/pyspark/ml/tests/test_linalg.py
@@ -15,17 +15,9 @@
 # limitations under the License.
 #
 
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
-
+import unittest
 import array as pyarray
+
 from numpy import arange, array, array_equal, inf, ones, tile, zeros
 
 from pyspark.ml.linalg import DenseMatrix, DenseVector, MatrixUDT, SparseMatrix, SparseVector, \
diff --git a/python/pyspark/ml/tests/test_param.py b/python/pyspark/ml/tests/test_param.py
@@ -19,17 +19,7 @@
 import inspect
 import sys
 import array as pyarray
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
-
-if sys.version > '3':
-    xrange = range
+import unittest
 
 import numpy as np
 
@@ -45,6 +35,10 @@
 from pyspark.testing.mlutils import check_params, PySparkTestCase, SparkSessionTestCase
 
 
+if sys.version > '3':
+    xrange = range
+
+
 class ParamTypeConversionTests(PySparkTestCase):
     """
     Test that param type conversion happens.
diff --git a/python/pyspark/ml/tests/test_persistence.py b/python/pyspark/ml/tests/test_persistence.py
@@ -17,16 +17,8 @@
 
 import json
 from shutil import rmtree
-import sys
 import tempfile
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 from pyspark.ml import Transformer
 from pyspark.ml.classification import DecisionTreeClassifier, LogisticRegression, OneVsRest, \
diff --git a/python/pyspark/ml/tests/test_pipeline.py b/python/pyspark/ml/tests/test_pipeline.py
@@ -14,15 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 from pyspark.ml.pipeline import Pipeline
 from pyspark.testing.mlutils import MockDataset, MockEstimator, MockTransformer, PySparkTestCase
diff --git a/python/pyspark/ml/tests/test_stat.py b/python/pyspark/ml/tests/test_stat.py
@@ -15,15 +15,7 @@
 # limitations under the License.
 #
 
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 from pyspark.ml.linalg import Vectors
 from pyspark.ml.stat import ChiSquareTest
diff --git a/python/pyspark/ml/tests/test_training_summary.py b/python/pyspark/ml/tests/test_training_summary.py
@@ -16,14 +16,7 @@
 #
 
 import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 if sys.version > '3':
     basestring = str
diff --git a/python/pyspark/ml/tests/test_tuning.py b/python/pyspark/ml/tests/test_tuning.py
@@ -15,16 +15,8 @@
 # limitations under the License.
 #
 
-import sys
 import tempfile
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 from pyspark.ml import Estimator, Model
 from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel, OneVsRest
diff --git a/python/pyspark/ml/tests/test_wrapper.py b/python/pyspark/ml/tests/test_wrapper.py
@@ -15,15 +15,7 @@
 # limitations under the License.
 #
 
-import sys
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 import py4j
 
diff --git a/python/pyspark/mllib/tests/test_algorithms.py b/python/pyspark/mllib/tests/test_algorithms.py
@@ -16,27 +16,16 @@
 #
 
 import os
-import sys
 import tempfile
 from shutil import rmtree
+import unittest
 
 from numpy import array, array_equal
-
 from py4j.protocol import Py4JJavaError
 
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
-
 from pyspark.mllib.fpm import FPGrowth
 from pyspark.mllib.recommendation import Rating
 from pyspark.mllib.regression import LabeledPoint
-from pyspark.sql.utils import IllegalArgumentException
 from pyspark.testing.mllibutils import make_serializer, MLlibTestCase
 
 
diff --git a/python/pyspark/mllib/tests/test_feature.py b/python/pyspark/mllib/tests/test_feature.py
@@ -15,20 +15,11 @@
 # limitations under the License.
 #
 
-import sys
 from math import sqrt
+import unittest
 
 from numpy import array, random, exp, abs, tile
 
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
-
 from pyspark.mllib.linalg import Vector, SparseVector, DenseVector, VectorUDT, Vectors
 from pyspark.mllib.linalg.distributed import RowMatrix
 from pyspark.mllib.feature import HashingTF, IDF, StandardScaler, ElementwiseProduct, Word2Vec
diff --git a/python/pyspark/mllib/tests/test_linalg.py b/python/pyspark/mllib/tests/test_linalg.py
@@ -17,18 +17,9 @@
 
 import sys
 import array as pyarray
+import unittest
 
 from numpy import array, array_equal, zeros, arange, tile, ones, inf
-from numpy import sum as array_sum
-
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
 
 import pyspark.ml.linalg as newlinalg
 from pyspark.mllib.linalg import Vector, SparseVector, DenseVector, VectorUDT, _convert_to_vector, \
diff --git a/python/pyspark/mllib/tests/test_stat.py b/python/pyspark/mllib/tests/test_stat.py
@@ -15,20 +15,11 @@
 # limitations under the License.
 #
 
-import sys
 import array as pyarray
+import unittest
 
 from numpy import array
 
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
-
 from pyspark.mllib.linalg import Vector, SparseVector, DenseVector, VectorUDT, _convert_to_vector, \
     DenseMatrix, SparseMatrix, Vectors, Matrices, MatrixUDT
 from pyspark.mllib.random import RandomRDDs
diff --git a/python/pyspark/mllib/tests/test_streaming_algorithms.py b/python/pyspark/mllib/tests/test_streaming_algorithms.py
@@ -15,21 +15,12 @@
 # limitations under the License.
 #
 
-import sys
 from time import time, sleep
+import unittest
 
 from numpy import array, random, exp, dot, all, mean, abs
 from numpy import sum as array_sum
 
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
-
 from pyspark import SparkContext
 from pyspark.mllib.clustering import StreamingKMeans, StreamingKMeansModel
 from pyspark.mllib.classification import StreamingLogisticRegressionWithSGD
diff --git a/python/pyspark/mllib/tests/test_util.py b/python/pyspark/mllib/tests/test_util.py
@@ -16,25 +16,14 @@
 #
 
 import os
-import sys
 import tempfile
-
-if sys.version_info[:2] <= (2, 6):
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')
-        sys.exit(1)
-else:
-    import unittest
+import unittest
 
 from pyspark.mllib.common import _to_java_object_rdd
 from pyspark.mllib.util import LinearDataGenerator
 from pyspark.mllib.util import MLUtils
-from pyspark.mllib.linalg import SparseVector, DenseVector, SparseMatrix, Vectors
+from pyspark.mllib.linalg import SparseVector, DenseVector, Vectors
 from pyspark.mllib.random import RandomRDDs
-from pyspark.mllib.regression import LabeledPoint
-from pyspark.mllib.stat import Statistics
 from pyspark.testing.mllibutils import MLlibTestCase
 
 
