Merge images after parallel write_doc_serialized

useblocks · Nov 3, 2023 · fd14612 · fd14612
1 parent f76523b
commit fd14612
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 9 deletions.
diff --git a/sphinx/builders/__init__.py b/sphinx/builders/__init__.py
@@ -603,12 +603,20 @@ def _write_serial(self, docnames: Sequence[str]) -> None:
                 self.write_doc(docname, doctree)
 
     def _write_parallel(self, docnames: Sequence[str], nproc: int) -> None:
+        # warm up caches/compile templates using the first document
+        firstname, docnames = docnames[0], docnames[1:]
+        self.app.phase = BuildPhase.RESOLVING
+        doctree = self.env.get_and_resolve_doctree(firstname, self)
+        self.app.phase = BuildPhase.WRITING
+        self.write_doc_serialized(firstname, doctree)
+        self.write_doc(firstname, doctree)
+
         def write_process(docs: list[tuple[str, nodes.document]]) -> bytes | None:
             self.app.phase = BuildPhase.WRITING
             for docname, doctree in docs:
                 if self.parallel_post_transform_ok:
                     doctree = self.env.get_and_resolve_doctree(docname, self, doctree)
-                self.write_doc_serialized(docname, doctree)
+                    self.write_doc_serialized(docname, doctree)
                 self.write_doc(docname, doctree)
             if self.parallel_post_transform_ok:
                 merge_attr = {
@@ -618,14 +626,6 @@ def write_process(docs: list[tuple[str, nodes.document]]) -> bytes | None:
                 return pickle.dumps(merge_attr, pickle.HIGHEST_PROTOCOL)
             return None
 
-        # warm up caches/compile templates using the first document
-        firstname, docnames = docnames[0], docnames[1:]
-        self.app.phase = BuildPhase.RESOLVING
-        doctree = self.env.get_and_resolve_doctree(firstname, self)
-        self.app.phase = BuildPhase.WRITING
-        self.write_doc_serialized(firstname, doctree)
-        self.write_doc(firstname, doctree)
-
         tasks = ParallelTasks(nproc)
         chunks = make_chunks(docnames, nproc)
 
@@ -648,6 +648,7 @@ def merge_builder(args: list[tuple[str, NoneType]], new_attrs_pickle: bytes) ->
                     doctree = self.env.get_doctree_write(docname)
                 else:
                     doctree = self.env.get_and_resolve_doctree(docname, self)
+                    self.write_doc_serialized(docname, doctree)
                 arg.append((docname, doctree))
             tasks.add_task(write_process, arg, merge_builder)
 

diff --git a/sphinx/builders/_epub_base.py b/sphinx/builders/_epub_base.py
@@ -155,6 +155,7 @@ class EpubBuilder(StandaloneHTMLBuilder):
     refuri_re = REFURI_RE
     template_dir = ""
     doctype = ""
+    post_transform_merge_attr = ['images']
 
     def init(self) -> None:
         super().init()
@@ -167,6 +168,17 @@ def init(self) -> None:
         self.use_index = self.get_builder_config('use_index', 'epub')
         self.refnodes: list[dict[str, Any]] = []
 
+    def merge_env_post_transform(self, new_attrs: dict[str, Any]) -> None:
+        """Merge images back to the main builder after parallel
+        post-transformation.
+
+        param new_attrs: the attributes from the parallel subprocess to be
+                         udpated in the main builder (self)
+        """
+        for filepath, filename in new_attrs['images'].items():
+            if filepath not in self.images:
+                self.images[filepath] = filename
+
     def create_build_info(self) -> BuildInfo:
         return BuildInfo(self.config, self.tags, ['html', 'epub'])
 

diff --git a/sphinx/builders/html/__init__.py b/sphinx/builders/html/__init__.py
@@ -181,6 +181,7 @@ class StandaloneHTMLBuilder(Builder):
 
     imgpath: str = ''
     domain_indices: list[DOMAIN_INDEX_TYPE] = []
+    post_transform_merge_attr = ['images']
 
     def __init__(self, app: Sphinx, env: BuildEnvironment) -> None:
         super().__init__(app, env)
@@ -233,6 +234,17 @@ def init(self) -> None:
 
         self.use_index = self.get_builder_config('use_index', 'html')
 
+    def merge_env_post_transform(self, new_attrs: dict[str, Any]) -> None:
+        """Merge images back to the main builder after parallel
+        post-transformation.
+
+        param new_attrs: the attributes from the parallel subprocess to be
+                         udpated in the main builder (self)
+        """
+        for filepath, filename in new_attrs['images'].items():
+            if filepath not in self.images:
+                self.images[filepath] = filename
+
     def create_build_info(self) -> BuildInfo:
         return BuildInfo(self.config, self.tags, ['html'])