DOP-1282: Make postprocessor nondestructive

Author: i80and

snooty/parser.py

@@ -1,5 +1,4 @@
 import collections
-import dataclasses
 import docutils.nodes
 import logging
 import multiprocessing
@@ -12,7 +11,17 @@
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path, PurePath
-from typing import Any, Dict, MutableSequence, Tuple, Optional, Set, List, Iterable
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    MutableSequence,
+    Tuple,
+    Optional,
+    Set,
+    List,
+    Iterable,
+)
 from docutils.nodes import make_id
 from typing_extensions import Protocol
 import docutils.utils
@@ -808,6 +817,62 @@ def on_delete(self, page_id: FileId, build_identifiers: BuildIdentifierSet) -> N
         ...
 
 
+class PageDatabase:
+    """A database of FileId->Page mappings that ensures the postprocessing pipeline
+       is run correctly. Raw parsed pages are added, flush() is called, then postprocessed
+       pages can be accessed."""
+
+    def __init__(self, postprocessor_factory: Callable[[], Postprocessor]) -> None:
+        self.postprocessor_factory = postprocessor_factory
+        self.parsed: Dict[FileId, Page] = {}
+        self.__postprocessed: Dict[FileId, Page] = {}
+        self.__changed_pages: Set[FileId] = set()
+
+    def __setitem__(self, key: FileId, value: Page) -> None:
+        """Set a raw parsed page."""
+        self.parsed[key] = value
+        self.__changed_pages.add(key)
+
+    def __getitem__(self, key: FileId) -> Page:
+        """If the postprocessor has been run since modifications were made, fetch a postprocessed page."""
+        assert not self.__changed_pages
+        return self.__postprocessed[key]
+
+    def __contains__(self, key: FileId) -> bool:
+        """Check if a given page exists in the parsed set."""
+        return key in self.parsed
+
+    def values(self) -> Iterable[Page]:
+        """Iterate over postprocessed pages."""
+        assert not self.__changed_pages
+        return self.__postprocessed.values()
+
+    def items(self) -> Iterable[Tuple[FileId, Page]]:
+        """Iterate over the postprocessed (FileId, Page) set."""
+        assert not self.__changed_pages
+        return self.__postprocessed.items()
+
+    def flush(
+        self
+    ) -> Tuple[Dict[str, SerializableType], Dict[FileId, List[Diagnostic]]]:
+        """Run the postprocessor if and only if any pages have changed, and return postprocessing results."""
+        if not self.__changed_pages:
+            return {}, {}
+
+        postprocessor = self.postprocessor_factory()
+
+        with util.PerformanceLogger.singleton().start("copy"):
+            copied_pages = util.fast_deep_copy(self.parsed)
+
+        with util.PerformanceLogger.singleton().start("postprocessing"):
+            post_metadata, post_diagnostics = postprocessor.run(copied_pages)
+
+        self.__postprocessed = postprocessor.pages
+        self.__changed_pages.clear()
+
+        return post_metadata, post_diagnostics
+
+
 class _Project:
     """Internal representation of a Snooty project with no data locking."""
 
@@ -832,12 +897,6 @@ def __init__(
         self.filesystem_watcher = filesystem_watcher
         self.build_identifiers = build_identifiers
 
-        self.postprocessor = (
-            DevhubPostprocessor(self.config, self.targets)
-            if self.config.default_domain == "devhub"
-            else Postprocessor(self.config, self.targets)
-        )
-
         self.yaml_mapping: Dict[str, GizaCategory[Any]] = {
             "steps": gizaparser.steps.GizaStepsCategory(self.config),
             "extracts": gizaparser.extracts.GizaExtractsCategory(self.config),
@@ -866,7 +925,11 @@ def __init__(
         ).strip()
         self.prefix = [self.config.name, username, branch]
 
-        self.pages: Dict[FileId, Page] = {}
+        self.pages = PageDatabase(
+            lambda: DevhubPostprocessor(self.config, self.targets)
+            if self.config.default_domain == "devhub"
+            else Postprocessor(self.config, self.targets)
+        )
 
         self.asset_dg: "networkx.DiGraph[FileId]" = networkx.DiGraph()
         self.expensive_operation_cache: Cache[FileId] = Cache()
@@ -906,11 +969,13 @@ def get_page_ast(self, path: Path) -> n.Node:
         """Update page file (.txt) with current text and return fully populated page AST"""
         # Get incomplete AST of page
         fileid = self.get_fileid(path)
+        post_metadata, post_diagnostics = self.pages.flush()
+        for fileid, diagnostics in post_diagnostics.items():
+            self.backend.on_diagnostics(fileid, diagnostics)
         page = self.pages[fileid]
 
-        # Fill in missing include nodes
         assert isinstance(page.ast, n.Parent)
-        return self._populate_include_nodes(page.ast)
+        return page.ast
 
     def get_project_name(self) -> str:
         return self.config.name
@@ -1042,8 +1107,7 @@ def create_page(filename: str) -> Tuple[Page, EmbeddedRstParser]:
                         page, all_yaml_diagnostics.get(page.source_path, [])
                     )
 
-        with util.PerformanceLogger.singleton().start("postprocessing"):
-            post_metadata, post_diagnostics = self.postprocessor.run(self.pages)
+        post_metadata, post_diagnostics = self.pages.flush()
 
         static_files = {
             "objects.inv": self.targets.generate_inventory("").dumps(
@@ -1052,7 +1116,7 @@ def create_page(filename: str) -> Tuple[Page, EmbeddedRstParser]:
         }
         post_metadata["static_files"] = static_files
 
-        for fileid, page in self.postprocessor.pages.items():
+        for fileid, page in self.pages.items():
             self.backend.on_update(self.prefix, self.build_identifiers, fileid, page)
         for fileid, diagnostics in post_diagnostics.items():
             self.backend.on_diagnostics(fileid, diagnostics)
@@ -1061,62 +1125,6 @@ def create_page(filename: str) -> Tuple[Page, EmbeddedRstParser]:
             self.prefix, self.build_identifiers, post_metadata
         )
 
-    def _populate_include_nodes(self, root: n.Parent[n.Node]) -> n.Node:
-        """
-        Add include nodes to page AST's children.
-
-        To render images on the Snooty extension's Snooty Preview,
-        we must use the full path of the image on the user's local machine. Note that this does change the
-        figure's value within the parser's dict. However, this should not change the value when using the parser
-        outside of Snooty Preview, since this function is currently only called by the language server.
-        """
-
-        def replace_nodes(node: n.Node) -> n.Node:
-            if isinstance(node, n.Directive):
-                if node.name == "include":
-                    # Get the name of the file
-                    argument = node.argument[0]
-                    include_filename = argument.value
-                    include_filename = include_filename[1:]
-
-                    # Get children of include file
-                    include_file_page_ast = self.pages[FileId(include_filename)].ast
-                    assert isinstance(include_file_page_ast, n.Parent)
-                    include_node_children = include_file_page_ast.children
-
-                    # Resolve includes within include node
-                    replaced_include = list(map(replace_nodes, include_node_children))
-                    node.children = replaced_include
-                # Replace instances of an image's name with its full path. This allows Snooty Preview to render an image by
-                # using the location of the image on the user's local machine
-                elif node.name == "figure":
-                    # Obtain subset of the image's path (name)
-                    argument = node.argument[0]
-                    image_value = argument.value
-
-                    # Prevents the image from having a redundant path if Snooty Preview already replaced
-                    # its original value.
-                    source_path_str = self.config.source_path.as_posix()
-                    index_match = image_value.find(source_path_str)
-                    if index_match != -1:
-                        repeated_offset = index_match + len(source_path_str)
-                        image_value = image_value[repeated_offset:]
-
-                    # Replace subset of path with full path of image
-                    if image_value[0] == "/":
-                        image_value = image_value[1:]
-                    full_path = self.get_full_path(FileId(image_value))
-                    argument.value = full_path.as_posix()
-            # Check for include nodes among current node's children
-            elif isinstance(node, n.Parent):
-                for child in node.children:
-                    replace_nodes(child)
-            return node
-
-        return dataclasses.replace(
-            root, children=list(map(replace_nodes, root.children))
-        )
-
     def _page_updated(self, page: Page, diagnostics: List[Diagnostic]) -> None:
         """Update any state associated with a parsed page."""
         # Finish any pending tasks
@@ -1130,7 +1138,7 @@ def _page_updated(self, page: Page, diagnostics: List[Diagnostic]) -> None:
         logger.debug("Updated: %s", fileid)
 
         if fileid in self.pages:
-            old_page = self.pages[fileid]
+            old_page = self.pages.parsed[fileid]
             old_assets = old_page.static_assets
             removed_assets = old_page.static_assets.difference(page.static_assets)
 

snooty/test_language_server.py

@@ -140,9 +140,6 @@ def test_text_doc_get_page_ast() -> None:
 
         source_path = server.project.config.source_path
 
-        # Image found in test file
-        image_path = Path("images/compass-create-database.png")
-        full_image_path = source_path.joinpath(image_path)
         # Change image path to be full path
         index_ast_string = (
             """<root>
@@ -153,7 +150,7 @@ def test_text_doc_get_page_ast() -> None:
             <heading id="id1"><text>Guides</text></heading>
             <directive name="figure" alt="Sample images" checksum="10e351828f156afcafc7744c30d7b2564c6efba1ca7c55cac59560c67581f947">
             <text>"""
-            + full_image_path.as_posix()
+            + "/images/compass-create-database.png"
             + """</text></directive>
             <directive name="include"><text>/includes/test_rst.rst</text>
             <directive name="include"><text>/includes/include_child.rst</text>

snooty/util.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import pickle
 import time
 import docutils.nodes
 import docutils.parsers.rst.directives
@@ -27,6 +28,7 @@
 from . import n
 
 logger = logging.getLogger(__name__)
+_T = TypeVar("_T")
 _K = TypeVar("_K", bound=Hashable)
 SOURCE_FILE_EXTENSIONS = {".txt", ".rst", ".yaml"}
 RST_EXTENSIONS = {".txt", ".rst"}
@@ -238,6 +240,12 @@ def split_domain(name: str) -> Tuple[str, str]:
     return parts[0], parts[1]
 
 
+def fast_deep_copy(d: Dict[_K, _T]) -> Dict[_K, _T]:
+    """Time-efficiently create deep copy of a dictionary containing trusted data.
+       This implementation currently invokes pickle, so should NOT be called on untrusted objects."""
+    return {k: pickle.loads(pickle.dumps(v)) for k, v in d.items()}
+
+
 class PerformanceLogger:
     _singleton: Optional["PerformanceLogger"] = None