1#!/usr/bin/python
2# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16#
17# Automatically copy TensorFlow binaries
18#
19# Usage:
20#           ./tensorflow/tools/ci_build/copy_binary.py --filename
21# tf_nightly/tf_nightly_gpu-1.4.0.dev20170914-cp35-cp35m-manylinux1_x86_64.whl
22# --new_py_ver 36
23#
24"""Copy binaries of TensorFlow for different python versions."""
25
26# pylint: disable=superfluous-parens
27
28import argparse
29import os
30import re
31import shutil
32import tempfile
33import zipfile
34
35TF_NIGHTLY_REGEX = (r"(.+)(tf_nightly.*)-(\d\.[\d]{1,2}"
36                    r"\.\d.dev[\d]{0,8})-(.+)\.whl")
37BINARY_STRING_TEMPLATE = "%s-%s-%s.whl"
38
39
40def check_existence(filename):
41  """Check the existence of file or dir."""
42  if not os.path.exists(filename):
43    raise RuntimeError("%s not found." % filename)
44
45
46def copy_binary(directory, origin_tag, new_tag, version, package):
47  """Rename and copy binaries for different python versions.
48
49  Args:
50    directory: string of directory
51    origin_tag: str of the old python version tag
52    new_tag: str of the new tag
53    version: the version of the package
54    package: str, name of the package
55
56  """
57  print("Rename and copy binaries with %s to %s." % (origin_tag, new_tag))
58  origin_binary = BINARY_STRING_TEMPLATE % (package, version, origin_tag)
59  new_binary = BINARY_STRING_TEMPLATE % (package, version, new_tag)
60  zip_ref = zipfile.ZipFile(os.path.join(directory, origin_binary), "r")
61
62  try:
63    tmpdir = tempfile.mkdtemp()
64    os.chdir(tmpdir)
65
66    zip_ref.extractall()
67    zip_ref.close()
68    old_py_ver = re.search(r"(cp\d\d-cp\d\d)", origin_tag).group(1)
69    new_py_ver = re.search(r"(cp\d\d-cp\d\d)", new_tag).group(1)
70
71    wheel_file = os.path.join(
72        tmpdir, "%s-%s.dist-info" % (package, version), "WHEEL")
73    with open(wheel_file, "r") as f:
74      content = f.read()
75    with open(wheel_file, "w") as f:
76      f.write(content.replace(old_py_ver, new_py_ver))
77
78    zout = zipfile.ZipFile(directory + new_binary, "w", zipfile.ZIP_DEFLATED)
79    zip_these_files = [
80        "%s-%s.dist-info" % (package, version),
81        "%s-%s.data" % (package, version),
82        "tensorflow",
83        "tensorflow_core",
84    ]
85    for dirname in zip_these_files:
86      for root, _, files in os.walk(dirname):
87        for filename in files:
88          zout.write(os.path.join(root, filename))
89    zout.close()
90  finally:
91    shutil.rmtree(tmpdir)
92
93
94def main():
95  """This script copies binaries.
96
97  Requirements:
98    filename: The path to the whl file
99    AND
100    new_py_ver: Create a nightly tag with current date
101
102  Raises:
103    RuntimeError: If the whl file was not found
104  """
105
106  parser = argparse.ArgumentParser(description="Cherry picking automation.")
107
108  # Arg information
109  parser.add_argument(
110      "--filename", help="path to whl file we are copying", required=True)
111  parser.add_argument(
112      "--new_py_ver", help="two digit py version eg. 27 or 33", required=True)
113
114  args = parser.parse_args()
115
116  # Argument checking
117  args.filename = os.path.abspath(args.filename)
118  check_existence(args.filename)
119  regex_groups = re.search(TF_NIGHTLY_REGEX, args.filename)
120  directory = regex_groups.group(1)
121  package = regex_groups.group(2)
122  version = regex_groups.group(3)
123  origin_tag = regex_groups.group(4)
124  old_py_ver = re.search(r"(cp\d\d)", origin_tag).group(1)
125
126  # Create new tags
127  new_tag = origin_tag.replace(old_py_ver, "cp" + args.new_py_ver)
128
129  # Copy the binary with the info we have
130  copy_binary(directory, origin_tag, new_tag, version, package)
131
132
133if __name__ == "__main__":
134  main()
135