def format_nb(
*,
notebooks: List[str],
remove_outputs: bool = False,
indent: int = 2,
test: bool = False,
) -> Status:
"""Formats a notebook."""
found_error = False # Track errors for final return code.
test_fail_notebooks = []
paths, err_paths = notebook_utils.collect_notebook_paths(notebooks)
if err_paths:
found_error = True
test_fail_notebooks.extend(err_paths)
for path in paths:
print(f"Format notebook: {path}", file=sys.stderr)
data, source = notebook_utils.load_notebook(path)
if not data:
found_error = True
test_fail_notebooks.append(path)
continue
# Returns formatted JSON byte string.
expected_output = clean_notebook(data, source, path, remove_outputs,
indent)
if test:
# Compare formatted contents with original file contents.
src_bytes = path.read_bytes()
if expected_output != src_bytes:
test_fail_notebooks.append(path)
else:
path.write_bytes(expected_output)
if test:
if test_fail_notebooks:
error_template = textwrap.dedent("""
[test] The following notebooks are not formatted:
{notebooks}
Please install `nbfmt` and format:
$ python3 -m pip install -U --user git+https://github.com/tensorflow/docs
$ python3 -m tensorflow_docs.tools.nbfmt notebook.ipynb
""")
notebooks = "\n".join(
[f"- {str(fp)}" for fp in test_fail_notebooks])
print(error_template.format(notebooks=notebooks), file=sys.stderr)
return Status.FAIL
else:
print("[test] Notebooks are formatted", file=sys.stderr)
return Status.PASS
if found_error:
return Status.FAIL
return Status.PASS
def _load_notebook(self, path: pathlib.Path) -> Tuple[Dict[str, Any], str]:
"""Load and parse JSON data from a notebook file.
Args:
path: A `pathlib.Path` of a Jupyter notebook.
Returns:
Dict: Contains data of the parsed JSON notebook.
String: The entire JSON source code of the notebook.
"""
data, source = notebook_utils.load_notebook(path)
if not data:
sys.exit(1) # load_notebook prints warning.
return data, source
def format_nb(
*,
notebooks: List[str],
remove_outputs: bool = False,
indent: int = 2,
test: bool = False,
) -> Status:
"""Formats a notebook."""
found_error = False # Track errors for final return code.
test_fail_notebooks = []
paths, err_paths = notebook_utils.collect_notebook_paths(notebooks)
if err_paths:
found_error = True
test_fail_notebooks.extend(err_paths)
for path in paths:
print(f"Format notebook: {path}", file=sys.stderr)
data, source = notebook_utils.load_notebook(path)
if not data:
found_error = True
test_fail_notebooks.append(path)
continue
clean_root(data, filepath=path) # Top-level notebook fields.
clean_cells(data, source, remove_outputs)
update_license_cells(data)
nbjson = json.dumps(data,
sort_keys=True,
ensure_ascii=False,
indent=indent)
if not OSS:
# Serialization differences in enviroments.
str_replaces = {"<": r"\u003c", ">": r"\u003e", "&": r"\u0026"}
for str_from, str_to in str_replaces.items():
nbjson = nbjson.replace(str_from, str_to)
expected_output = (nbjson + "\n").encode("utf-8")
if test:
# Compare formatted contents with original file contents.
src_bytes = path.read_bytes()
if expected_output != src_bytes:
test_fail_notebooks.append(path)
else:
path.write_bytes(expected_output)
if test:
if test_fail_notebooks:
error_template = textwrap.dedent("""
[test] The following notebooks are not formatted:
{notebooks}
Please install `nbfmt` and format:
$ python3 -m pip install -U --user git+https://github.com/tensorflow/docs
$ python3 -m tensorflow_docs.tools.nbfmt notebook.ipynb
""")
notebooks = "\n".join(
[f"- {str(fp)}" for fp in test_fail_notebooks])
print(error_template.format(notebooks=notebooks), file=sys.stderr)
return Status.FAIL
else:
print("[test] Notebooks are formatted", file=sys.stderr)
return Status.PASS
if found_error:
return Status.FAIL
return Status.PASS
def main(argv):
if len(argv) <= 1:
raise app.UsageError("Missing arguments.")
found_error = False # Track errors for final return code.
test_fail_notebooks = []
paths, err_paths = notebook_utils.collect_notebook_paths(argv[1:])
if err_paths:
found_error = True
test_fail_notebooks.extend(err_paths)
for path in paths:
print(f"Format notebook: {path}", file=sys.stderr)
data, source = notebook_utils.load_notebook(path)
if not data:
found_error = True
test_fail_notebooks.append(path)
continue
# Set top-level notebook defaults.
data["nbformat"] = 4
data["nbformat_minor"] = 0
remove_extra_fields(data) # Top-level fields.
clean_cells(data, source, FLAGS.remove_outputs)
update_metadata(data, filepath=path)
update_license_cell(data)
nbjson = json.dumps(data,
sort_keys=True,
ensure_ascii=False,
indent=FLAGS.indent)
# Some enviroments require a different format.
if not OSS:
nbjson = nbjson.replace("<", r"\u003c").replace(">", r"\u003e")
expected_output = (nbjson + "\n").encode("utf-8")
if FLAGS.test:
# Compare formatted contents with original file contents.
src_bytes = path.read_bytes()
if expected_output != src_bytes:
test_fail_notebooks.append(path)
else:
path.write_bytes(expected_output)
if FLAGS.test:
if test_fail_notebooks:
error_template = textwrap.dedent("""
[test] The following notebooks are not formatted:
{notebooks}
Please install `nbfmt` and format:
$ python3 -m pip install -U --user git+https://github.com/tensorflow/docs
$ python3 -m tensorflow_docs.tools.nbfmt notebook.ipynb
""")
notebooks = "\n".join(
[f"- {str(fp)}" for fp in test_fail_notebooks])
print(error_template.format(notebooks=notebooks), file=sys.stderr)
sys.exit(1)
else:
print("[test] Notebooks are formatted", file=sys.stderr)
sys.exit(0)
if found_error:
sys.exit(1)