EVOLUTION-MANAGER

Edit File: ipynb.py

# Lint as: python2, python3
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""A module to support operations on ipynb files"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import copy
import json
import re
import shutil
import tempfile

import six

CodeLine = collections.namedtuple("CodeLine", ["cell_number", "code"])

def is_python(cell):
  """Checks if the cell consists of Python code."""
  return (cell["cell_type"] == "code"  # code cells only
          and cell["source"]  # non-empty cells
          and not six.ensure_str(cell["source"][0]).startswith("%%")
         )  # multiline eg: %%bash

def process_file(in_filename, out_filename, upgrader):
  """The function where we inject the support for ipynb upgrade."""
  print("Extracting code lines from original notebook")
  raw_code, notebook = _get_code(in_filename)
  raw_lines = [cl.code for cl in raw_code]

# The function follows the original flow from `upgrader.process_fil`
  with tempfile.NamedTemporaryFile("w", delete=False) as temp_file:

processed_file, new_file_content, log, process_errors = (
        upgrader.update_string_pasta("\n".join(raw_lines), in_filename))

if temp_file and processed_file:
      new_notebook = _update_notebook(
          notebook, raw_code,
          six.ensure_str(new_file_content).split("\n"))
      json.dump(new_notebook, temp_file)
    else:
      raise SyntaxError(
          "Was not able to process the file: \n%s\n" % "".join(log))

files_processed = processed_file
    report_text = upgrader._format_log(log, in_filename, out_filename)
    errors = process_errors

shutil.move(temp_file.name, out_filename)

return files_processed, report_text, errors

def skip_magic(code_line, magic_list):
  """Checks if the cell has magic, that is not Python-based.

Args:
      code_line: A line of Python code
      magic_list: A list of jupyter "magic" exceptions

Returns:
    If the line jupyter "magic" line, not Python line

>>> skip_magic('!ls -laF', ['%', '!', '?'])
  True
  """

for magic in magic_list:
    if six.ensure_str(code_line).startswith(magic):
      return True

return False

def check_line_split(code_line):
  r"""Checks if a line was split with `\`.

Args:
      code_line: A line of Python code

Returns:
    If the line was split with `\`

>>> skip_magic("!gcloud ml-engine models create ${MODEL} \\\n")
  True
  """

return re.search(r"\\\s*\n$", code_line)

def _get_code(input_file):
  """Loads the ipynb file and returns a list of CodeLines."""

raw_code = []

with open(input_file) as in_file:
    notebook = json.load(in_file)

cell_index = 0
  for cell in notebook["cells"]:
    if is_python(cell):
      cell_lines = cell["source"]

is_line_split = False
      for line_idx, code_line in enumerate(cell_lines):

# Sometimes, jupyter has more than python code
        # Idea is to comment these lines, for upgrade time
        if skip_magic(code_line, ["%", "!", "?"]) or is_line_split:
          # Found a special character, need to "encode"
          code_line = "###!!!" + six.ensure_str(code_line)

# if this cell ends with `\` -> skip the next line
          is_line_split = check_line_split(code_line)

if is_line_split:
          is_line_split = check_line_split(code_line)

# Sometimes, people leave \n at the end of cell
        # in order to migrate only related things, and make the diff
        # the smallest -> here is another hack
        if (line_idx == len(cell_lines) -
            1) and six.ensure_str(code_line).endswith("\n"):
          code_line = six.ensure_str(code_line).replace("\n", "###===")

# sometimes a line would start with `\n` and content after
        # that's the hack for this
        raw_code.append(
            CodeLine(cell_index,
                     six.ensure_str(code_line.rstrip()).replace("\n",
                                                                "###===")))

cell_index += 1

return raw_code, notebook

def _update_notebook(original_notebook, original_raw_lines, updated_code_lines):
  """Updates notebook, once migration is done."""

new_notebook = copy.deepcopy(original_notebook)

# validate that the number of lines is the same
  assert len(original_raw_lines) == len(updated_code_lines), \
    ("The lengths of input and converted files are not the same: "
     "{} vs {}".format(len(original_raw_lines), len(updated_code_lines)))

code_cell_idx = 0
  for cell in new_notebook["cells"]:
    if not is_python(cell):
      continue

applicable_lines = [
        idx for idx, code_line in enumerate(original_raw_lines)
        if code_line.cell_number == code_cell_idx
    ]

new_code = [updated_code_lines[idx] for idx in applicable_lines]

cell["source"] = "\n".join(new_code).replace("###!!!", "").replace(
        "###===", "\n")
    code_cell_idx += 1

return new_notebook