Simplify dev_tools.notebooks.utils.rewrite_notebook

dev_tools/notebooks/__init__.py

@@ -12,4 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dev_tools.notebooks.utils import filter_notebooks, list_all_notebooks, rewrite_notebook
+from dev_tools.notebooks.utils import (
+    filter_notebooks,
+    list_all_notebooks,
+    remove_if_temporary_notebook,
+    rewrite_notebook,
+)

dev_tools/notebooks/isolated_notebook_test.py

@@ -33,7 +33,12 @@
 import pytest
 
 from dev_tools import shell_tools
-from dev_tools.notebooks import list_all_notebooks, filter_notebooks, rewrite_notebook
+from dev_tools.notebooks import (
+    list_all_notebooks,
+    filter_notebooks,
+    rewrite_notebook,
+    remove_if_temporary_notebook,
+)
 
 # these notebooks rely on features that are not released yet
 # after every release we should raise a PR and empty out this list
@@ -177,7 +182,7 @@ def test_notebooks_against_released_cirq(partition, notebook_path, cloned_env):
 
     notebook_file = os.path.basename(notebook_path)
 
-    rewritten_notebook_descriptor, rewritten_notebook_path = rewrite_notebook(notebook_path)
+    rewritten_notebook_path = rewrite_notebook(notebook_path)
 
     cmd = f"""
 mkdir -p out/{notebook_rel_dir}
@@ -208,9 +213,7 @@ def test_notebooks_against_released_cirq(partition, notebook_path, cloned_env):
             f"instead of `pip install cirq` to this notebook, and exclude it from "
             f"dev_tools/notebooks/isolated_notebook_test.py."
         )
-
-    if rewritten_notebook_descriptor:
-        os.close(rewritten_notebook_descriptor)
+    remove_if_temporary_notebook(rewritten_notebook_path)
 
 
 @pytest.mark.parametrize("notebook_path", NOTEBOOKS_DEPENDING_ON_UNRELEASED_FEATURES)

dev_tools/notebooks/notebook_test.py

@@ -24,7 +24,12 @@
 import pytest
 
 from dev_tools import shell_tools
-from dev_tools.notebooks import filter_notebooks, list_all_notebooks, rewrite_notebook
+from dev_tools.notebooks import (
+    filter_notebooks,
+    list_all_notebooks,
+    rewrite_notebook,
+    remove_if_temporary_notebook,
+)
 
 SKIP_NOTEBOOKS = [
     # skipping vendor notebooks as we don't have auth sorted out
@@ -67,7 +72,7 @@ def test_notebooks_against_released_cirq(notebook_path):
     notebook_file = os.path.basename(notebook_path)
     notebook_rel_dir = os.path.dirname(os.path.relpath(notebook_path, "."))
     out_path = f"out/{notebook_rel_dir}/{notebook_file[:-6]}.out.ipynb"
-    rewritten_notebook_descriptor, rewritten_notebook_path = rewrite_notebook(notebook_path)
+    rewritten_notebook_path = rewrite_notebook(notebook_path)
     papermill_flags = "--no-request-save-on-cell-execute --autosave-cell-every 0"
     cmd = f"""mkdir -p out/{notebook_rel_dir}
 papermill {rewritten_notebook_path} {out_path} {papermill_flags}"""
@@ -83,6 +88,4 @@ def test_notebooks_against_released_cirq(notebook_path):
             f"notebook (in Github Actions, you can download it from the workflow artifact"
             f" 'notebook-outputs')"
         )
-
-    if rewritten_notebook_descriptor:
-        os.close(rewritten_notebook_descriptor)
+    remove_if_temporary_notebook(rewritten_notebook_path)

dev_tools/notebooks/utils.py

@@ -69,7 +69,7 @@ def rewrite_notebook(notebook_path):
 
         * Lines in this file without `->` are ignored.
 
-        * Lines in this file with `->` are split into two (if there are mulitple `->` it is an
+        * Lines in this file with `->` are split into two (if there are multiple `->` it is an
         error). The first of these is compiled into a pattern match, via `re.compile`, and
         the second is the replacement for that match.
 
@@ -82,16 +82,17 @@ def rewrite_notebook(notebook_path):
     It is the responsibility of the caller of this method to delete the new file.
 
     Returns:
-        Tuple of a file descriptor and the file path for the rewritten file.  If no `.tst` file
-        was found, then the file descriptor is None and the path is `notebook_path`.
+        The file path for the rewritten file.  If no `.tst` file was found,
+        return the input `notebook_path` as the notebook was not changed.
+        Otherwise return a new path created in the temporary directory.
 
     Raises:
         AssertionError: If there are multiple `->` per line, or not all of the replacements
             are used.
     """
     notebook_test_path = os.path.splitext(notebook_path)[0] + '.tst'
     if not os.path.exists(notebook_test_path):
-        return None, notebook_path
+        return notebook_path
 
     # Get the rewrite rules.
     patterns = []
@@ -104,20 +105,37 @@ def rewrite_notebook(notebook_path):
 
     used_patterns = set()
     with open(notebook_path, 'r') as original_file:
-        new_file_descriptor, new_file_path = tempfile.mkstemp(suffix='.ipynb')
-        with open(new_file_path, 'w') as new_file:
-            for line in original_file:
-                new_line = line
-                for pattern, replacement in patterns:
-                    new_line = pattern.sub(replacement, new_line)
-                    if new_line != line:
-                        used_patterns.add(pattern)
-                        break
-                new_file.write(new_line)
+        lines = original_file.readlines()
+    for i, line in enumerate(lines):
+        for pattern, replacement in patterns:
+            new_line = pattern.sub(replacement, line)
+            if new_line != line:
+                lines[i] = new_line
+                used_patterns.add(pattern)
+                break
 
     assert len(patterns) == len(used_patterns), (
         'Not all patterns where used. Patterns not used: '
         f'{set(x for x, _ in patterns) - used_patterns}'
     )
 
-    return new_file_descriptor, new_file_path
+    with tempfile.NamedTemporaryFile(mode='w', suffix='-rewrite.ipynb', delete=False) as new_file:
+        new_file.writelines(lines)
+
+    return new_file.name
+
+
+def remove_if_temporary_notebook(notebook_path: str):
+    """Delete temporary notebook if written by `rewrite_notebook`.
+
+    Do nothing if the specified notebook is not in the temporary directory
+    or if its filename does not end in '-rewrite.ipynb'
+
+    Use this to safely clean up notebooks created by `rewrite_notebook`.
+    """
+    if (
+        notebook_path.endswith('-rewrite.ipynb')
+        and os.path.isfile(notebook_path)
+        and os.path.dirname(notebook_path) == tempfile.gettempdir()
+    ):
+        os.remove(notebook_path)

dev_tools/notebooks/utils_test.py

@@ -37,40 +37,43 @@ def write_test_data(ipynb_txt, tst_txt):
 def test_rewrite_notebook():
     directory, ipynb_path = write_test_data('d = 5\nd = 4', 'd = 5->d = 3')
 
-    descriptor, path = dt.rewrite_notebook(ipynb_path)
+    path = dt.rewrite_notebook(ipynb_path)
 
     assert path != ipynb_path
     with open(path, 'r') as f:
         rewritten = f.read()
         assert rewritten == 'd = 3\nd = 4'
 
-    os.close(descriptor)
+    dt.remove_if_temporary_notebook(path)
+    assert not os.path.exists(path)
     shutil.rmtree(directory)
 
 
 def test_rewrite_notebook_multiple():
     directory, ipynb_path = write_test_data('d = 5\nd = 4', 'd = 5->d = 3\nd = 4->d = 1')
 
-    descriptor, path = dt.rewrite_notebook(ipynb_path)
+    path = dt.rewrite_notebook(ipynb_path)
 
     with open(path, 'r') as f:
         rewritten = f.read()
         assert rewritten == 'd = 3\nd = 1'
 
-    os.close(descriptor)
+    dt.remove_if_temporary_notebook(path)
+    assert not os.path.exists(path)
     shutil.rmtree(directory)
 
 
 def test_rewrite_notebook_ignore_non_seperator_lines():
     directory, ipynb_path = write_test_data('d = 5\nd = 4', 'd = 5->d = 3\n# comment')
 
-    descriptor, path = dt.rewrite_notebook(ipynb_path)
+    path = dt.rewrite_notebook(ipynb_path)
 
     with open(path, 'r') as f:
         rewritten = f.read()
         assert rewritten == 'd = 3\nd = 4'
 
-    os.close(descriptor)
+    dt.remove_if_temporary_notebook(path)
+    assert not os.path.exists(path)
     shutil.rmtree(directory)
 
 
@@ -80,10 +83,12 @@ def test_rewrite_notebook_no_tst_file():
     with open(ipynb_path, 'w') as f:
         f.write('d = 5\nd = 4')
 
-    descriptor, path = dt.rewrite_notebook(ipynb_path)
+    path = dt.rewrite_notebook(ipynb_path)
 
-    assert descriptor is None
+    # verify clean up is skipped when files are not rewritten
     assert path == ipynb_path
+    dt.remove_if_temporary_notebook(path)
+    assert os.path.exists(path)
 
     shutil.rmtree(directory)