Cut call time of add_prefix() by 99.8%

src/prefixmaps/datamodel/context.py

@@ -5,7 +5,7 @@
 from collections import defaultdict
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import List, Mapping, Optional
+from typing import List, Mapping, Optional, Set, Union
 
 import curies
 
@@ -143,6 +143,14 @@ class Context:
     merged_from: Optional[List[str]] = None
     upper: bool = None
     lower: bool = None
+    _prefixes: Set[str] = field(default_factory=set)
+    """Private attr to speed up duplicate lookups"""
+    _prefixes_lower: Set[str] = field(default_factory=set)
+    """Private attr to speed up duplicate lookups"""
+    _namespaces: Set[str] = field(default_factory=set)
+    """Private attr to speed up duplicate lookups"""
+    _namespaces_lower: Set[str] = field(default_factory=set)
+    """Private attr to speed up duplicate lookups"""
 
     def combine(self, context: "Context"):
         """
@@ -164,6 +172,7 @@ def add_prefix(
         status: StatusType = StatusType.canonical,
         preferred: bool = False,
         expansion_source: Optional[str] = None,
+        force: bool = False,
     ):
         """
         Adds a prefix expansion to this context.
@@ -182,18 +191,20 @@ def add_prefix(
         :param expansion_source: An optional annotation to be used when merging contexts together.
             The source will keep track of the original context that a given prefix
             expansion came from. This is used in :meth:`Context.combine`.
+        :param force: if True, recompute namespaces and prefixes. default False.
         :return:
         """
         # TODO: check status
+        _prefix = prefix
         if not preferred:
             if self.upper:
                 prefix = prefix.upper()
                 if self.lower:
                     raise ValueError("Cannot set both upper AND lower")
             if self.lower:
                 prefix = prefix.lower()
-        prefixes = self.prefixes(lower=True)
-        namespaces = self.namespaces(lower=True)
+        prefixes = self.prefixes(lower=True, force=force, as_list=False)
+        namespaces = self.namespaces(lower=True, force=force, as_list=False)
         if prefix.lower() in prefixes:
             if namespace.lower() in namespaces:
                 return
@@ -203,6 +214,7 @@ def add_prefix(
         else:
             if namespace.lower() in namespaces:
                 status = StatusType.namespace_alias
+
         self.prefix_expansions.append(
             PrefixExpansion(
                 context=self.name,
@@ -212,6 +224,10 @@ def add_prefix(
                 expansion_source=expansion_source,
             )
         )
+        self._prefixes.add(_prefix)
+        self._prefixes_lower.add(prefix.lower())
+        self._namespaces.add(namespace)
+        self._namespaces_lower.add(namespace.lower())
 
     def filter(self, prefix: PREFIX = None, namespace: NAMESPACE = None):
         """
@@ -230,29 +246,57 @@ def filter(self, prefix: PREFIX = None, namespace: NAMESPACE = None):
             filtered_pes.append(pe)
         return filtered_pes
 
-    def prefixes(self, lower=False) -> List[str]:
+    def prefixes(
+        self, lower=False, force: bool = True, as_list: bool = True
+    ) -> Union[List[str], Set[str]]:
         """
         All unique prefixes in all prefix expansions.
 
         :param lower: if True, the prefix is normalized to lowercase.
+        :param force: if True, recompute. if False, return cached
+        :param as_list: if True (default), return as a list. Otherwise a set
         :return:
         """
         if lower:
-            return list({pe.prefix.lower() for pe in self.prefix_expansions})
+            if force or len(self._prefixes_lower) == 0:
+                self._prefixes_lower = {pe.prefix.lower() for pe in self.prefix_expansions}
+            res = self._prefixes_lower
+
+        else:
+            if force or len(self._prefixes) == 0:
+                self._prefixes = {pe.prefix for pe in self.prefix_expansions}
+            res = self._prefixes
+
+        if as_list:
+            return list(res)
         else:
-            return list({pe.prefix for pe in self.prefix_expansions})
+            return res
 
-    def namespaces(self, lower=False) -> List[str]:
+    def namespaces(
+        self, lower=False, force: bool = True, as_list: bool = True
+    ) -> Union[List[str], Set[str]]:
         """
         All unique namespaces in all prefix expansions
 
         :param lower: if True, the namespace is normalized to lowercase.
+        :param force: if True, recompute. if False, return cached
+        :param as_list: if True (default), return as a list. Otherwise a set
         :return:
         """
         if lower:
-            return list({pe.namespace.lower() for pe in self.prefix_expansions})
+            if force or len(self._namespaces_lower) == 0:
+                self._namespaces_lower = {pe.namespace.lower() for pe in self.prefix_expansions}
+            res = self._namespaces_lower
+
+        else:
+            if force or len(self._namespaces) == 0:
+                self._namespaces = {pe.namespace for pe in self.prefix_expansions}
+            res = self._namespaces
+
+        if as_list:
+            return list(res)
         else:
-            return list({pe.namespace for pe in self.prefix_expansions})
+            return res
 
     def as_dict(self) -> PREFIX_EXPANSION_DICT:
         """