Explicit python version dependency
diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index 6b90eca..9771a5c 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -66,6 +66,13 @@
'tensorflow_estimator >= 1.14.0rc0, < 1.15.0rc0',
'termcolor >= 1.1.0',
'wrapt >= 1.11.1',
+ # python3 requires wheel 0.26
+ 'wheel >= 0.26;python_version>="3"',
+ 'wheel;python_version<"3"',
+ # mock comes with unittest.mock for python3, need to install for python2
+ 'mock >= 2.0.0;python_version<"3"',
+ # functools comes with python3, need to install the backport for python2
+ 'functools32 >= 3.2.3;python_version<"3"',
]
if sys.byteorder == 'little':
@@ -81,16 +88,6 @@
sys.argv.remove('--project_name')
sys.argv.pop(project_name_idx)
-# python3 requires wheel 0.26
-if sys.version_info.major == 3:
- REQUIRED_PACKAGES.append('wheel >= 0.26')
-else:
- REQUIRED_PACKAGES.append('wheel')
- # mock comes with unittest.mock for python3, need to install for python2
- REQUIRED_PACKAGES.append('mock >= 2.0.0')
- # functools comes with python3, need to install the backport for python2
- REQUIRED_PACKAGES.append('functools32 >= 3.2.3')
-
# tf-nightly should depend on tb-nightly
if 'tf_nightly' in project_name:
for i, pkg in enumerate(REQUIRED_PACKAGES):
