tools: Update notebook script

tools/export_to_ipynb.py

@@ -2,25 +2,87 @@
 """Transform any Python sample or example to Python NoteBook."""
 import ast
 import sys
+import os
 from nbformat import v3
 from nbformat import v4
 
 input_file = sys.argv[1]
-print('reading %s' % input_file)
+print(f'reading {input_file}')
 with open(input_file) as fpin:
   text = fpin.read()
 
+# Compute output file path.
+output_file = input_file
+output_file = output_file.replace('.py', '.ipynb')
+# For example/python/foo.py -> example/notebook/examples/foo.ipynb
+output_file = output_file.replace('examples/python',
+                                  'examples/notebook/examples')
+# For example/contrib/foo.py -> example/notebook/contrib/foo.ipynb
+output_file = output_file.replace('examples/contrib',
+                                  'examples/notebook/contrib')
+# For ortools/*/samples/foo.py -> example/notebook/*/foo.ipynb
+output_file = output_file.replace('ortools', 'examples/notebook')
+output_file = output_file.replace('samples/', '')
+
+
 nbook = v3.reads_py('')
 nbook = v4.upgrade(nbook)  # Upgrade v3 to v4
 
+print(f'Adding copyright cell...')
+copyright = '##### Copyright 2020 The OR-Tools Authors.'
+nbook['cells'].append(v4.new_markdown_cell(copyright))
+
+print(f'Adding license cell...')
+license = '''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.
+'''
+nbook['cells'].append(v4.new_markdown_cell(license))
+
+print(f'Adding Title cell...')
+basename = '# ' + os.path.basename(input_file).replace('.py', '')
+nbook['cells'].append(v4.new_markdown_cell(basename))
+
+print(f'Adding link cell...')
+github_logo = 'https://raw.githubusercontent.com/google/or-tools/master/tools/github_32px.png'
+github_path = 'https://github.com/google/or-tools/blob/master/' + input_file
+
+colab_path =  'https://colab.research.google.com/github/google/or-tools/blob/master/' + output_file
+colab_logo = 'https://raw.githubusercontent.com/google/or-tools/master/tools/colab_32px.png'
+link = f'''<table align=\"left\">
+<td>
+<a href=\"{colab_path}\"><img src=\"{colab_logo}\"/>Run in Google Colab</a>
+</td>
+<td>
+<a href=\"{github_path}\"><img src=\"{github_logo}\"/>View source on GitHub</a>
+</td>
+</table>'''
+nbook['cells'].append(v4.new_markdown_cell(link))
+
+print(f'Installing ortools cell...')
+install_doc = 'First, you must install [ortools](https://pypi.org/project/ortools/) package in this colab.'
+nbook['cells'].append(v4.new_markdown_cell(install_doc))
+install_cmd = '!pip install ortools'
+nbook['cells'].append(v4.new_code_cell(install_cmd))
+
+print(f'Adding code cell...')
 all_blocks = ast.parse(text).body
+print(f'number of bocks: {len(all_blocks)}')
 line_start = [c.lineno - 1 for c in all_blocks]
 line_start[0] = 0
 lines = text.split('\n')
 
 full_text = ''
-
 for c_block, s, e in zip(all_blocks, line_start, line_start[1:] + [len(lines)]):
+  print(c_block)
   c_text = '\n'.join(lines[s:e])
   if isinstance(c_block,
                 ast.If) and c_block.test.comparators[0].s == '__main__':
@@ -44,20 +106,7 @@ for c_block, s, e in zip(all_blocks, line_start, line_start[1:] + [len(lines)]):
 nbook['cells'].append(v4.new_code_cell(full_text))
 
 jsonform = v4.writes(nbook) + '\n'
-output_file = input_file
-output_file = output_file.replace('.py', '.ipynb')
-# For example/python/foo.py -> example/notebook/examples/foo.ipynb
-output_file = output_file.replace('examples/python',
-                                  'examples/notebook/examples')
 
-# For example/contrib/foo.py -> example/notebook/contrib/foo.ipynb
-output_file = output_file.replace('examples/contrib',
-                                  'examples/notebook/contrib')
-
-# For ortools/*/samples/foo.py -> example/notebook/*/foo.ipynb
-output_file = output_file.replace('ortools', 'examples/notebook')
-output_file = output_file.replace('samples/', '')
-
-print('writing %s' % output_file)
+print(f'writing {output_file}')
 with open(output_file, 'w') as fpout:
   fpout.write(jsonform)