From 62d12da2e331854fa0160e13d9a1e3948f1381ad Mon Sep 17 00:00:00 2001
From: Sergii Khomenko <sergii.khomenko@stylight.com>
Date: Sun, 14 Apr 2019 15:12:57 +0200
Subject: [PATCH] Add support for line splits

---
 tensorflow/tools/compatibility/ipynb.py | 40 +++++++++++++++++++++----
 1 file changed, 35 insertions(+), 5 deletions(-)

diff --git a/tensorflow/tools/compatibility/ipynb.py b/tensorflow/tools/compatibility/ipynb.py
index d37a1abde25..d2656e1321c 100644
--- a/tensorflow/tools/compatibility/ipynb.py
+++ b/tensorflow/tools/compatibility/ipynb.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =============================================================================
-"""A module to support operation on ipynb files"""
+"""A module to support operations on ipynb files"""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -62,8 +62,36 @@ def process_file(in_filename, out_filename, upgrader):
   return files_processed, report_text, errors
 
 
+def skip_magic(code_line, magic_list):
+    """
+    Checks if the cell has magic, that is not python based
+
+    >>> skip_magic('!ls -laF')
+    True
+    """
+
+    for magic in magic_list:
+        if code_line.startswith(magic):
+            return True
+
+    return False
+
+
+def check_line_split(code_line):
+    r"""
+    Checks if line was splitted with `\`
+    >>> skip_magic("!gcloud ml-engine models create ${MODEL} \\\n")
+    True
+    """
+
+    if code_line.endswith('\\\n'):
+        return True
+
+    return False
+
+
 def _get_code(input_file):
-  """Load the ipynb file and return a list of CodeLines."""
+  """Loads the ipynb file and returns a list of CodeLines."""
 
   raw_code = []
 
@@ -75,15 +103,17 @@ def _get_code(input_file):
     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 code_line.startswith("%") or code_line.startswith("!") \
-            or code_line.startswith("?"):
+        if skip_magic(code_line, ['%', '!', '?']) or is_line_split:
           # Found a special character, need to "encode"
           code_line = "###!!!" + code_line
 
+        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
@@ -102,7 +132,7 @@ def _get_code(input_file):
 
 
 def _update_notebook(original_notebook, original_raw_lines, updated_code_lines):
-  """Update notebook, once migration is done."""
+  """Updates notebook, once migration is done."""
 
   new_notebook = copy.deepcopy(original_notebook)