🔧 Improve internal API for needs access (#1255)

This PR makes it clearer when we are accessing the (mutable) unprocessed needs, during the read/collection phase,
or when we are accessing the (immutable) processed needs, during the write/analysis phase.
It does not affect the user facing API.

It also adds the abstracted `NeedsView` type for the processed needs.
This is not changed from the simple dict (`id -> need dict`), but
it provides the groundwork for potential future changes to make this more performant for large amounts of needs.

docs/api.rst

@@ -30,3 +30,9 @@ Exceptions
 ----------
 .. automodule:: sphinx_needs.api.exceptions
    :members:
+
+Data
+----
+
+.. automodule:: sphinx_needs.data
+   :members: NeedsInfoType, NeedsView

docs/conf.py

@@ -49,7 +49,6 @@
     ("py:class", "docutils.statemachine.StringList"),
     ("py:class", "T"),
     ("py:class", "sphinx_needs.debug.T"),
-    ("py:class", "sphinx_needs.data.NeedsInfoType"),
 ]
 
 rst_epilog = """
@@ -757,7 +756,7 @@ def create_tutorial_needs(app: Sphinx, _env, _docnames):
 
     We do this dynamically, to avoid having to maintain the JSON file manually.
     """
-    all_data = SphinxNeedsData(app.env).get_or_create_needs()
+    all_data = SphinxNeedsData(app.env).get_needs_view()
     writer = NeedsList(app.config, outdir=app.confdir, confdir=app.confdir)
     for i in range(1, 5):
         test_id = f"T_00{i}"

sphinx_needs/api/need.py

@@ -312,7 +312,7 @@ def run():
     # Add need to global need list
     #############################################################################################
 
-    if need_id in SphinxNeedsData(env).get_or_create_needs():
+    if SphinxNeedsData(env).has_need(need_id):
         if id:
             message = f"A need with ID {need_id} already exists, " f"title: {title!r}."
         else:  # this is a generated ID
@@ -457,7 +457,7 @@ def run():
     if needs_info["post_template"]:
         needs_info["post_content"] = _prepare_template(app, needs_info, "post_template")
 
-    SphinxNeedsData(env).get_or_create_needs()[need_id] = needs_info
+    SphinxNeedsData(env).add_need(needs_info)
 
     if needs_info["is_external"]:
         return []
@@ -604,12 +604,10 @@ def del_need(app: Sphinx, need_id: str) -> None:
     :param need_id: Sphinx need id.
     """
     data = SphinxNeedsData(app.env)
-    needs = data.get_or_create_needs()
-    if need_id in needs:
-        del needs[need_id]
-        data.remove_need_node(need_id)
-    else:
+    if not data.has_need(need_id):
         log_warning(logger, f"Given need id {need_id} not exists!", None, None)
+    else:
+        data.remove_need(need_id)
 
 
 def add_external_need(

sphinx_needs/builder.py

@@ -85,7 +85,7 @@ def finish(self) -> None:
 
         filter_string = needs_config.builder_filter
         filtered_needs: list[NeedsInfoType] = filter_needs(
-            data.get_or_create_needs().values(),
+            data.get_needs_view().values(),
             needs_config,
             filter_string,
             append_warning="(from need_builder_filter)",
@@ -178,7 +178,7 @@ def finish(self) -> None:
 
         data = SphinxNeedsData(self.env)
         needs = (
-            data.get_or_create_needs().values()
+            data.get_needs_view().values()
         )  # We need a list of needs for later filter checks
         version = getattr(self.env.config, "version", "unset")
         needs_config = NeedsSphinxConfig(self.env.config)

sphinx_needs/data.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Final, Literal, Mapping, TypedDict
+from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Mapping, NewType, TypedDict
 
 from sphinx.util.logging import getLogger
 
@@ -676,23 +676,85 @@ class NeedsUmlType(NeedsBaseDataType):
     content_calculated: str
 
 
+NeedsMutable = NewType("NeedsMutable", Dict[str, NeedsInfoType])
+"""A mutable view of the needs, before resolution
+"""
+
+NeedsView = NewType("NeedsView", Mapping[str, NeedsInfoType])
+"""A read-only view of the needs, after resolution
+(e.g. back links have been computed etc)
+"""
+
+
 class SphinxNeedsData:
     """Centralised access to sphinx-needs data, stored within the Sphinx environment."""
 
     def __init__(self, env: BuildEnvironment) -> None:
         self.env = env
 
-    def get_or_create_needs(self) -> dict[str, NeedsInfoType]:
-        """Get all needs, mapped by ID.
-
-        This is lazily created and cached in the environment.
-        """
+    @property
+    def _env_needs(self) -> dict[str, NeedsInfoType]:
         try:
             return self.env.needs_all_needs
         except AttributeError:
             self.env.needs_all_needs = {}
         return self.env.needs_all_needs
 
+    def has_need(self, need_id: str) -> bool:
+        """Check if a need with the given ID exists."""
+        return need_id in self._env_needs
+
+    def add_need(self, need: NeedsInfoType) -> None:
+        """Add an unprocessed need to the cache.
+
+        This will overwrite any existing need with the same ID.
+
+        .. important:: this should only be called within the read phase,
+            before the needs have been fully collected and resolved.
+        """
+        self._env_needs[need["id"]] = need
+
+    def remove_need(self, need_id: str) -> None:
+        """Remove a single need from the cache, if it exists.
+
+        .. important:: this should only be called within the read phase,
+            before the needs have been fully collected and resolved.
+        """
+        if need_id in self._env_needs:
+            del self._env_needs[need_id]
+        self.remove_need_node(need_id)
+
+    def remove_doc(self, docname: str) -> None:
+        """Remove all data related to a document from the cache.
+
+        .. important:: this should only be called within the read phase,
+            before the needs have been fully collected and resolved.
+        """
+        for need_id in list(self._env_needs):
+            if self._env_needs[need_id]["docname"] == docname:
+                del self._env_needs[need_id]
+                self.remove_need_node(need_id)
+        docs = self.get_or_create_docs()
+        for key, value in docs.items():
+            docs[key] = [doc for doc in value if doc != docname]
+
+    def get_needs_mutable(self) -> NeedsMutable:
+        """Get all needs, mapped by ID.
+
+        .. important:: this should only be called within the read phase,
+            before the needs have been fully collected and resolved.
+        """
+        return self._env_needs  # type: ignore[return-value]
+
+    def get_needs_view(self) -> NeedsView:
+        """Return a read-only view of all needs, after resolution.
+
+        .. important:: this should only be called within the write phase,
+            after the needs have been fully collected
+            and resolved (e.g. back links have been computed etc)
+        """
+        return self._env_needs  # type: ignore[return-value]
+
     @property
     def has_export_filters(self) -> bool:
         """Whether any filters have export IDs."""
@@ -793,7 +855,7 @@ def remove_need_node(self, need_id: str) -> None:
             del self._needs_all_nodes[need_id]
 
     def get_need_node(self, need_id: str) -> Need | None:
-        """Get a need node from the cache, if it exists."""
+        """Get a copy of a need node from the cache, if it exists."""
         if need_id in self._needs_all_nodes:
             # We must create a copy of the node, as it may be reused several time
             # (multiple needextract for the same need) and the Sphinx ImageTransformator add location specific
@@ -822,8 +884,8 @@ def merge_data(
     this_data.get_or_create_filters().update(other_data.get_or_create_filters())
 
     # Update needs
-    needs = this_data.get_or_create_needs()
-    other_needs = other_data.get_or_create_needs()
+    needs = this_data._env_needs
+    other_needs = other_data._env_needs
     for other_id, other_need in other_needs.items():
         if other_id in needs:
             # we only want to warn if the need comes from one of the docs parsed in this worker

sphinx_needs/directives/need.py

@@ -15,7 +15,7 @@
 from sphinx_needs.api import add_need
 from sphinx_needs.api.exceptions import NeedsInvalidException
 from sphinx_needs.config import NEEDS_CONFIG, NeedsSphinxConfig
-from sphinx_needs.data import NeedsInfoType, SphinxNeedsData
+from sphinx_needs.data import NeedsMutable, SphinxNeedsData
 from sphinx_needs.debug import measure_time
 from sphinx_needs.defaults import NEED_DEFAULT_OPTIONS
 from sphinx_needs.directives.needextend import Needextend, extend_needs_data
@@ -284,12 +284,7 @@ def purge_needs(app: Sphinx, env: BuildEnvironment, docname: str) -> None:
     Gets executed, if a doc file needs to be purged/ read in again.
     So this code delete all found needs for the given docname.
     """
-    data = SphinxNeedsData(env)
-    needs = data.get_or_create_needs()
-    for need_id in list(needs):
-        if needs[need_id]["docname"] == docname:
-            del needs[need_id]
-            data.remove_need_node(need_id)
+    SphinxNeedsData(env).remove_doc(docname)
 
 
 def analyse_need_locations(app: Sphinx, doctree: nodes.document) -> None:
@@ -305,7 +300,7 @@ def analyse_need_locations(app: Sphinx, doctree: nodes.document) -> None:
     """
     env = app.env
 
-    needs = SphinxNeedsData(env).get_or_create_needs()
+    needs = SphinxNeedsData(env).get_needs_mutable()
 
     hidden_needs: list[Need] = []
     for need_node in doctree.findall(Need):
@@ -380,7 +375,7 @@ def post_process_needs_data(app: Sphinx) -> None:
     """
     needs_config = NeedsSphinxConfig(app.config)
     needs_data = SphinxNeedsData(app.env)
-    needs = needs_data.get_or_create_needs()
+    needs = needs_data.get_needs_mutable()
     if needs and not needs_data.needs_is_post_processed:
         extend_needs_data(needs, needs_data.get_or_create_extends(), needs_config)
         resolve_dynamic_values(needs, app)
@@ -406,7 +401,7 @@ def process_need_nodes(app: Sphinx, doctree: nodes.document, fromdocname: str) -
     needs_data = SphinxNeedsData(app.env)
 
     # If no needs were defined, we do not need to do anything
-    if not needs_data.get_or_create_needs():
+    if not needs_data.get_needs_view():
         return
 
     post_process_needs_data(app)
@@ -426,7 +421,7 @@ def format_need_nodes(
 ) -> None:
     """Replace need nodes in the document with node trees suitable for output"""
     env = app.env
-    needs = SphinxNeedsData(env).get_or_create_needs()
+    needs = SphinxNeedsData(env).get_needs_view()
 
     # We try to avoid findall as much as possibles. so we reuse the already found need nodes in the current document.
     # for node_need in doctree.findall(Need):
@@ -448,7 +443,7 @@ def format_need_nodes(
         node_need.parent.replace(node_need, rendered_node)
 
 
-def check_links(needs: dict[str, NeedsInfoType], config: NeedsSphinxConfig) -> None:
+def check_links(needs: NeedsMutable, config: NeedsSphinxConfig) -> None:
     """Checks if set links are valid or are dead (referenced need does not exist.)
 
     For needs with dead links, an extra ``has_dead_links`` field is added and,
@@ -493,9 +488,7 @@ def check_links(needs: dict[str, NeedsInfoType], config: NeedsSphinxConfig) -> N
                                 )
 
 
-def create_back_links(
-    needs: dict[str, NeedsInfoType], config: NeedsSphinxConfig
-) -> None:
+def create_back_links(needs: NeedsMutable, config: NeedsSphinxConfig) -> None:
     """Create back-links in all found needs.
 
     These are fields for each link type, ``<link_name>_back``,

sphinx_needs/directives/needbar.py

@@ -292,7 +292,7 @@ def process_needbar(
         # 5. process content
         local_data_number = []
         need_list = list(
-            prepare_need_list(needs_data.get_or_create_needs().values())
+            prepare_need_list(needs_data.get_needs_view().values())
         )  # adds parts to need_list
 
         for line in local_data:

sphinx_needs/directives/needextend.py

@@ -9,7 +9,11 @@
 
 from sphinx_needs.api.exceptions import NeedsInvalidFilter
 from sphinx_needs.config import NeedsSphinxConfig
-from sphinx_needs.data import NeedsExtendType, NeedsInfoType, SphinxNeedsData
+from sphinx_needs.data import (
+    NeedsExtendType,
+    NeedsMutable,
+    SphinxNeedsData,
+)
 from sphinx_needs.filter_common import filter_needs
 from sphinx_needs.logging import get_logger, log_warning
 from sphinx_needs.utils import add_doc
@@ -94,7 +98,7 @@ def _split_value(value: str) -> Sequence[tuple[str, bool]]:
 
 
 def extend_needs_data(
-    all_needs: dict[str, NeedsInfoType],
+    all_needs: NeedsMutable,
     extends: dict[str, NeedsExtendType],
     needs_config: NeedsSphinxConfig,
 ) -> None:
@@ -105,7 +109,7 @@ def extend_needs_data(
 
     for current_needextend in extends.values():
         need_filter = current_needextend["filter"]
-        if need_filter in all_needs:
+        if need_filter and need_filter in all_needs:
             # a single known ID
             found_needs = [all_needs[need_filter]]
         elif need_filter is not None and re.fullmatch(

sphinx_needs/directives/needextract.py

@@ -92,7 +92,7 @@ def process_needextract(
             continue
 
         current_needextract: NeedsExtractType = node.attributes
-        all_needs = SphinxNeedsData(env).get_or_create_needs()
+        all_needs = SphinxNeedsData(env).get_needs_view()
         content = nodes.container()
         content.attributes["ids"] = [current_needextract["target_id"]]
 

sphinx_needs/directives/needfilter.py

@@ -85,7 +85,7 @@ def process_needfilters(
     builder = app.builder
     env = app.env
     needs_config = NeedsSphinxConfig(env.config)
-    all_needs = SphinxNeedsData(env).get_or_create_needs()
+    all_needs = SphinxNeedsData(env).get_needs_view()
 
     # NEEDFILTER
     # for node in doctree.findall(Needfilter):

sphinx_needs/directives/needflow/_graphviz.py

@@ -49,7 +49,7 @@ def process_needflow_graphviz(
 ) -> None:
     needs_config = NeedsSphinxConfig(app.config)
     env_data = SphinxNeedsData(app.env)
-    all_needs = env_data.get_or_create_needs()
+    all_needs = env_data.get_needs_view()
 
     link_type_names = [link["option"].upper() for link in needs_config.extra_links]
     allowed_link_types_options = [link.upper() for link in needs_config.flow_link_types]

sphinx_needs/directives/needflow/_plantuml.py

@@ -203,7 +203,7 @@ def process_needflow_plantuml(
     env = app.env
     needs_config = NeedsSphinxConfig(app.config)
     env_data = SphinxNeedsData(env)
-    all_needs = env_data.get_or_create_needs()
+    all_needs = env_data.get_needs_view()
 
     link_type_names = [link["option"].upper() for link in needs_config.extra_links]
     allowed_link_types_options = [link.upper() for link in needs_config.flow_link_types]

sphinx_needs/directives/needflow/_shared.py

@@ -8,24 +8,24 @@
 from sphinx_needs.data import (
     NeedsFlowType,
     NeedsInfoType,
+    NeedsView,
 )
 from sphinx_needs.logging import get_logger
 
 logger = get_logger(__name__)
 
 
 def filter_by_tree(
-    all_needs: dict[str, NeedsInfoType],
+    all_needs: NeedsView,
     root_id: str,
     link_types: list[LinkOptionsType],
     direction: Literal["both", "incoming", "outgoing"],
     depth: int | None,
-) -> dict[str, NeedsInfoType]:
+) -> NeedsView:
     """Filter all needs by the given ``root_id``,
     and all needs that are connected to the root need by the given ``link_types``, in the given ``direction``."""
-    need_items: dict[str, NeedsInfoType] = {}
     if root_id not in all_needs:
-        return need_items
+        return NeedsView({})
     roots = {root_id: (0, all_needs[root_id])}
     link_prefixes = (
         ("_back",)
@@ -37,6 +37,7 @@ def filter_by_tree(
     links_to_process = [
         link["option"] + d for link in link_types for d in link_prefixes
     ]
+    need_items: dict[str, NeedsInfoType] = {}
     while roots:
         root_id, (root_depth, root) = roots.popitem()
         if root_id in need_items:
@@ -53,7 +54,7 @@ def filter_by_tree(
                 }
             )
 
-    return need_items
+    return NeedsView(need_items)
 
 
 def get_root_needs(found_needs: list[NeedsInfoType]) -> list[NeedsInfoType]: