Skip to content

Commit

Permalink
tools: format export_to_ipynb.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@Mizux
Mizux committed Aug 12, 2024
1 parent c460256 commit 419a8a0
Showing 1 changed file with 141 additions and 116 deletions.
257 changes: 141 additions & 116 deletions tools/export_to_ipynb.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,32 +22,30 @@
from nbformat import v4

input_file = sys.argv[1]
print(f'reading {input_file}')
print(f"reading {input_file}")
with open(input_file) as fpin:
text = fpin.read()
text = fpin.read()

# Compute output file path.
output_file = input_file
output_file = output_file.replace('.py', '.ipynb')
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')
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')
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/', '')
output_file = output_file.replace("ortools", "examples/notebook")
output_file = output_file.replace("samples/", "")

nbook = v3.reads_py('')
nbook = v3.reads_py("")
nbook = v4.upgrade(nbook) # Upgrade v3 to v4

print('Adding copyright cell...')
google = '##### Copyright 2023 Google LLC.'
nbook['cells'].append(v4.new_markdown_cell(source=google, id='google'))
print("Adding copyright cell...")
google = "##### Copyright 2024 Google LLC."
nbook["cells"].append(v4.new_markdown_cell(source=google, id="google"))

print('Adding license cell...')
apache = '''Licensed under the Apache License, Version 2.0 (the "License");
print("Adding license cell...")
apache = """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
Expand All @@ -58,115 +56,142 @@
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(source=apache, id='apache'))

print('Adding Title cell...')
basename = '# ' + os.path.basename(input_file).replace('.py', '')
nbook['cells'].append(v4.new_markdown_cell(source=basename, id='basename'))

print('Adding link cell...')
github_logo = 'https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png'
github_path = 'https://github.com/google/or-tools/blob/main/' + input_file

colab_path = 'https://colab.research.google.com/github/google/or-tools/blob/main/' + output_file
colab_logo = 'https://raw.githubusercontent.com/google/or-tools/main/tools/colab_32px.png'
link = f'''<table align=\"left\">
"""
nbook["cells"].append(v4.new_markdown_cell(source=apache, id="apache"))

print("Adding Title cell...")
basename = "# " + os.path.basename(input_file).replace(".py", "")
nbook["cells"].append(v4.new_markdown_cell(source=basename, id="basename"))

print("Adding link cell...")
github_logo = (
"https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png"
)
github_path = "https://github.com/google/or-tools/blob/main/" + input_file

colab_path = (
"https://colab.research.google.com/github/google/or-tools/blob/main/" + output_file
)
colab_logo = (
"https://raw.githubusercontent.com/google/or-tools/main/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(source=link, id='link'))

print('Adding ortools install cell...')
install_doc = ('First, you must install '
'[ortools](https://pypi.org/project/ortools/) package in this '
'colab.')
nbook['cells'].append(v4.new_markdown_cell(source=install_doc, id='doc'))
install_cmd = '%pip install ortools'
nbook['cells'].append(v4.new_code_cell(source=install_cmd, id='install'))

print('Adding code cell...')
</table>"""
nbook["cells"].append(v4.new_markdown_cell(source=link, id="link"))

print("Adding ortools install cell...")
install_doc = (
"First, you must install "
"[ortools](https://pypi.org/project/ortools/) package in this "
"colab."
)
nbook["cells"].append(v4.new_markdown_cell(source=install_doc, id="doc"))
install_cmd = "%pip install ortools"
nbook["cells"].append(v4.new_code_cell(source=install_cmd, id="install"))

print("Adding code cell...")
all_blocks = ast.parse(text).body
print(f'number of blocks: {len(all_blocks)}')
print(f"number of blocks: {len(all_blocks)}")
line_start = [c.lineno - 1 for c in all_blocks]
line_start[0] = 0
lines = text.split('\n')
lines = text.split("\n")

full_text = ''
full_text = ""
for idx, (c_block, s, e) in enumerate(
zip(all_blocks, line_start, line_start[1:] + [len(lines)])):
print(f'block[{idx}]: {c_block}')
c_text = '\n'.join(lines[s:e])
# Clean boilerplate header and description
if (idx == 0 and isinstance(c_block, ast.Expr) and
isinstance(c_block.value, ast.Constant)):
print('Adding description cell...')
filtered_lines = lines[s:e]
# filtered_lines = list(
# filter(lambda l: not l.startswith('#!'), lines[s:e]))
filtered_lines = list(
filter(lambda l: not re.search(r'^#!', l), filtered_lines))
filtered_lines = list(
filter(lambda l: not re.search(r'# \[START .*\]$', l), filtered_lines))
filtered_lines = list(
filter(lambda l: not re.search(r'# \[END .*\]$', l), filtered_lines))
# TODO(user): Remove only copyright not all line with '^#'
filtered_lines = list(
filter(lambda l: not l.startswith(r'#'), filtered_lines))
filtered_lines = [s.replace(r'"""', '') for s in filtered_lines]
filtered_text = '\n'.join(filtered_lines)
nbook['cells'].append(
v4.new_markdown_cell(source=filtered_text, id='description'))
# Remove absl app import
elif (isinstance(c_block, ast.ImportFrom) and c_block.module == 'absl'
and c_block.names[0].name == 'app'):
print(f'Removing import {c_block.module}.{c_block.names[0].name}...')
# rewrite absl flag import
elif (isinstance(c_block, ast.ImportFrom) and c_block.module == 'absl'
and c_block.names[0].name == 'flags'):
print(f'Rewrite import {c_block.module}.{c_block.names[0].name}...')
full_text += 'from ortools.sat.colab import flags\n'
# Unwrap __main__ function
elif (isinstance(c_block, ast.If) and
c_block.test.comparators[0].s == '__main__'):
print('Unwrapping main function...')
c_lines = lines[s + 1:e]
# remove start and de-indent lines
spaces_to_delete = c_block.body[0].col_offset
fixed_lines = [
n_line[spaces_to_delete:]
if n_line.startswith(' ' * spaces_to_delete) else n_line
for n_line in c_lines
]
filtered_lines = fixed_lines
filtered_lines = list(
filter(lambda l: not re.search(r'# \[START .*\]$', l), filtered_lines))
filtered_lines = list(
filter(lambda l: not re.search(r'# \[END .*\]$', l), filtered_lines))
filtered_lines = [
re.sub(r'app.run\((.*)\)$', r'\1()', s) for s in filtered_lines
]
full_text += '\n'.join(filtered_lines) + '\n'
# Others
else:
print('Appending block...')
filtered_lines = lines[s:e]
for i, line in enumerate(filtered_lines):
filtered_lines[i] = line.replace('DEFINE_', 'define_')
filtered_lines = list(
filter(lambda l: not re.search(r'# \[START .*\]$', l), filtered_lines))
filtered_lines = list(
filter(lambda l: not re.search(r'# \[END .*\]$', l), filtered_lines))
full_text += '\n'.join(filtered_lines) + '\n'

nbook['cells'].append(v4.new_code_cell(source=full_text, id='code'))

jsonform = v4.writes(nbook) + '\n'

print(f'writing {output_file}')
with open(output_file, 'w') as fpout:
fpout.write(jsonform)
zip(all_blocks, line_start, line_start[1:] + [len(lines)])
):
print(f"block[{idx}]: {c_block}")
c_text = "\n".join(lines[s:e])
# Clean boilerplate header and description
if (
idx == 0
and isinstance(c_block, ast.Expr)
and isinstance(c_block.value, ast.Constant)
):
print("Adding description cell...")
filtered_lines = lines[s:e]
# filtered_lines = list(
# filter(lambda l: not l.startswith('#!'), lines[s:e]))
filtered_lines = list(
filter(lambda l: not re.search(r"^#!", l), filtered_lines)
)
filtered_lines = list(
filter(lambda l: not re.search(r"# \[START .*\]$", l), filtered_lines)
)
filtered_lines = list(
filter(lambda l: not re.search(r"# \[END .*\]$", l), filtered_lines)
)
# TODO(user): Remove only copyright not all line with '^#'
filtered_lines = list(filter(lambda l: not l.startswith(r"#"), filtered_lines))
filtered_lines = [s.replace(r'"""', "") for s in filtered_lines]
filtered_text = "\n".join(filtered_lines)
nbook["cells"].append(
v4.new_markdown_cell(source=filtered_text, id="description")
)
# Remove absl app import
elif (
isinstance(c_block, ast.ImportFrom)
and c_block.module == "absl"
and c_block.names[0].name == "app"
):
print(f"Removing import {c_block.module}.{c_block.names[0].name}...")
# rewrite absl flag import
elif (
isinstance(c_block, ast.ImportFrom)
and c_block.module == "absl"
and c_block.names[0].name == "flags"
):
print(f"Rewrite import {c_block.module}.{c_block.names[0].name}...")
full_text += "from ortools.sat.colab import flags\n"
# Unwrap __main__ function
elif isinstance(c_block, ast.If) and c_block.test.comparators[0].s == "__main__":
print("Unwrapping main function...")
c_lines = lines[s + 1 : e]
# remove start and de-indent lines
spaces_to_delete = c_block.body[0].col_offset
fixed_lines = [
(
n_line[spaces_to_delete:]
if n_line.startswith(" " * spaces_to_delete)
else n_line
)
for n_line in c_lines
]
filtered_lines = fixed_lines
filtered_lines = list(
filter(lambda l: not re.search(r"# \[START .*\]$", l), filtered_lines)
)
filtered_lines = list(
filter(lambda l: not re.search(r"# \[END .*\]$", l), filtered_lines)
)
filtered_lines = [
re.sub(r"app.run\((.*)\)$", r"\1()", s) for s in filtered_lines
]
full_text += "\n".join(filtered_lines) + "\n"
# Others
else:
print("Appending block...")
filtered_lines = lines[s:e]
for i, line in enumerate(filtered_lines):
filtered_lines[i] = line.replace("DEFINE_", "define_")
filtered_lines = list(
filter(lambda l: not re.search(r"# \[START .*\]$", l), filtered_lines)
)
filtered_lines = list(
filter(lambda l: not re.search(r"# \[END .*\]$", l), filtered_lines)
)
full_text += "\n".join(filtered_lines) + "\n"

nbook["cells"].append(v4.new_code_cell(source=full_text, id="code"))

jsonform = v4.writes(nbook) + "\n"

print(f"writing {output_file}")
with open(output_file, "w") as fpout:
fpout.write(jsonform)

0 comments on commit 419a8a0

Please sign in to comment.