Cache functionality

ocrd_models/ocrd_models/ocrd_mets.py

@@ -4,6 +4,7 @@
 from datetime import datetime
 import re
 from lxml import etree as ET
+from copy import deepcopy
 
 from ocrd_utils import (
     is_local_filename,
@@ -41,7 +42,7 @@ class OcrdMets(OcrdXmlDocument):
     """
 
     @staticmethod
-    def empty_mets(now=None):
+    def empty_mets(now=None, cache_flag=False):
         """
         Create an empty METS file from bundled template.
         """
@@ -50,20 +51,91 @@ def empty_mets(now=None):
         tpl = METS_XML_EMPTY.decode('utf-8')
         tpl = tpl.replace('{{ VERSION }}', VERSION)
         tpl = tpl.replace('{{ NOW }}', '%s' % now)
-        return OcrdMets(content=tpl.encode('utf-8'))
+        return OcrdMets(content=tpl.encode('utf-8'), cache_flag=cache_flag)
 
     def __init__(self, **kwargs):
         """
 
         """
         super(OcrdMets, self).__init__(**kwargs)
 
+        # If cache is enabled
+        if self._cache_flag:
+            # Cache for the fileGrps (mets:fileGrp) - a dictionary with Key and Value pair:
+            # Key: 'fileGrp.USE'
+            # Value: a 'fileGrp' object at some memory location 
+            self._fileGrp_cache = {}
+
+            # Cache for the files (mets:file) - two nested dictionaries
+            # The outer dictionary's Key: 'fileGrp.USE'
+            # The outer dictionary's Value: Inner dictionary
+            # The inner dictionary's Key: 'file.ID'
+            # The inner dictionary's Value: a 'file' object at some memory location
+            self._file_cache = {}
+
+            # Note, if the empty_mets() function is used to instantiate OcrdMets
+            # Then the cache is empty even after this operation
+            self._fill_caches()
+
+    def __exit__(self):
+        if self._cache_flag:
+            self._clear_caches()
+
     def __str__(self):
         """
         String representation
         """
         return 'OcrdMets[fileGrps=%s,files=%s]' % (self.file_groups, list(self.find_files()))
 
+    def _fill_caches(self):
+        """
+        Fills the caches with fileGrps and FileIDs
+        """
+
+        tree_root = self._tree.getroot()
+        el_fileGrp_list = tree_root.find(".//mets:fileSec", NS)
+        if el_fileGrp_list is None or len(el_fileGrp_list) == 0:
+            return
+
+        for el_fileGrp in el_fileGrp_list:
+            fileGrp_use = el_fileGrp.get('USE')
+
+            # NOTE: For some reason the el_fileGrp_list contains None values
+            # when testing with the SBB0000F29300010000/data/mets.xml
+            if fileGrp_use is None:
+                continue
+
+            self._fileGrp_cache[fileGrp_use] = el_fileGrp
+            print("_fill_caches> file group added to the cache: %s" % fileGrp_use)
+
+            # Assign an empty dictionary that will hold the files of the added fileGrp
+            self._file_cache[fileGrp_use] = {}
+
+            for el_file in el_fileGrp:
+                file_id = el_file.get('ID')
+                self._file_cache[fileGrp_use].update({file_id : el_file})
+                print("_fill_caches> file added to the cache: %s" % file_id)
+
+        print("_fill_caches> total fileGrp cache elements: %s" % len(self._fileGrp_cache))
+
+    def _clear_caches(self):
+        """
+        Deallocates the caches
+        """
+
+        fileGrp_counter = 0
+
+        for key in list(self._fileGrp_cache):
+            del self._fileGrp_cache[key]
+            fileGrp_counter += 1
+
+        # print("_clear_caches> total cleared fileGrp cache elements: %d" % fileGrp_counter)
+
+        for key in list(self._file_cache):
+            for inner_key in list(self._file_cache[key]):
+                del self._file_cache[key][inner_key]
+            del self._file_cache[key]
+
     @property
     def unique_identifier(self):
         """
@@ -127,6 +199,95 @@ def find_all_files(self, *args, **kwargs):
 
         Equivalent to ``list(self.find_files(...))``
         """
+
+        # NOTE: This code gets complex with the REGEX.
+        # Having two separate funcitons: with REGEX and without REGEX would simplify things
+        if self._cache_flag:
+            matches = []
+
+            # If only both the fileGrp and ID parameters have been passed
+            # Faster search in the cache
+            if 'ID' in kwargs and 'fileGrp' in kwargs and 'pageId' not in kwargs and 'mimetype' not in kwargs and 'url' not in kwargs:
+                fileGrp = kwargs['fileGrp']
+                fileID = kwargs['ID']
+
+                if fileID.startswith(REGEX_PREFIX):
+                    fileID = re.compile(fileID[REGEX_PREFIX_LEN:])
+                if fileGrp.startswith(REGEX_PREFIX):
+                    fileGrp = re.compile(fileGrp[REGEX_PREFIX_LEN:])
+
+                # Case where no regex pattern is given and
+                # exact match could be obtained
+                if (isinstance(fileID, str) and isinstance(fileGrp, str)):
+                    if fileGrp in self._file_cache:  
+                        if fileID in self._file_cache[fileGrp]:
+                            matches.append(OcrdFile(self._file_cache[fileGrp][fileID], mets=self))
+                elif isinstance(fileGrp, str):
+                    # fileGrp is str and fileID is regex
+                    if fileGrp in self._file_cache:
+                        for fileID_str in self._file_cache[fileGrp]:
+                            if fileID.fullmatch(fileID_str):
+                                matches.append(OcrdFile(self._file_cache[fileGrp_str][fileID_str], mets=self))
+
+                elif isinstance(fileID, str):
+                    # fileID is str and fileGrp is regex
+                    for fileGrp_str in self._file_cache:
+                        if fileGrp.fullmatch(fileGrp_str):
+                            if fileID in self._file_cache[fileGrp_str]:
+                                matches.append(OcrdFile(self._file_cache[fileGrp_str][fileID], mets=self))
+
+                else:
+                    # both are regex: this has a really bad performance since
+                    # we have to iterate all groups and all files to check for matches
+                    for fileGrp_str in self._file_cache:
+                        if fileGrp.fullmatch(fileGrp_str):
+                            for fileID_str in self._file_cache[fileGrp_str]:
+                                if fileID.fullmatch(fileID_str):
+                                    matches.append(OcrdFile(self._file_cache[fileGrp_str][fileID_str], mets=self))
+
+                return matches
+
+            # If only the fileGrp parameter has been passed
+            # Return a list with all files of that fileGrp from the cache
+            if 'fileGrp' in kwargs and 'ID' not in kwargs and 'pageId' not in kwargs and 'mimetype' not in kwargs and 'url' not in kwargs:
+                fileGrp = kwargs['fileGrp']
+
+                if fileGrp.startswith(REGEX_PREFIX):
+                    fileGrp = re.compile(fileGrp[REGEX_PREFIX_LEN:])
+
+                    for fileGrp_str in self._file_cache:
+                        print("Type(fileGrp): %s, Type(fileGrp_str): %s" % (type(fileGrp), type(fileGrp_str)))
+                        if fileGrp.fullmatch(fileGrp_str):
+                            for fileID in self._file_cache[fileGrp_str]:
+                                matches.append(OcrdFile(self._file_cache[fileGrp_str][fileID], mets=self))
+                else:
+                    if fileGrp in self._file_cache:
+                        for fileID in self._file_cache[fileGrp]:
+                            matches.append(OcrdFile(self._file_cache[fileGrp][fileID], mets=self))
+
+                return matches
+
+            # If only the ID parameter has been passed
+            # Return a list with that fileID inside or an empty list if not in cache
+            if 'ID' in kwargs and 'fileGrp' not in kwargs and 'pageId' not in kwargs and 'mimetype' not in kwargs and 'url' not in kwargs:
+                fileID = kwargs['ID']
+
+                if fileID.startswith(REGEX_PREFIX):
+                    fileID = re.compile(fileID[REGEX_PREFIX_LEN:])
+                    for fileGrp_str in self._file_cache:
+                        for fileID_str in self._file_cache[fileGrp_str]:
+                            if fileID.fullmatch(fileID_str):
+                                matches.append(OcrdFile(self._file_cache[fileGrp_str][fileID_str], mets=self))
+
+                else:
+                    for fileGrp_str in self._file_cache:
+                        if fileID in self._file_cache[fileGrp_str]:
+                            matches.append(OcrdFile(self._file_cache[fileGrp_str][fileID], mets=self))
+
+                return matches
+
+        # Run the old routine if cache is not enabled
+        # or search based on parameters other than fileGrp and fileID are used
         return list(self.find_files(*args, **kwargs))
 
     # pylint: disable=multiple-statements
@@ -183,6 +344,7 @@ def find_files(self, ID=None, fileGrp=None, pageId=None, mimetype=None, url=None
             mimetype = re.compile(mimetype[REGEX_PREFIX_LEN:])
         if url and url.startswith(REGEX_PREFIX):
             url = re.compile(url[REGEX_PREFIX_LEN:])
+
         for cand in self._tree.getroot().xpath('//mets:file', namespaces=NS):
             if ID:
                 if isinstance(ID, str):
@@ -238,6 +400,13 @@ def add_file_group(self, fileGrp):
         if el_fileGrp is None:
             el_fileGrp = ET.SubElement(el_fileSec, TAG_METS_FILEGRP)
             el_fileGrp.set('USE', fileGrp)
+
+            # Add the fileGrp to both caches
+            if self._cache_flag:
+                self._fileGrp_cache[fileGrp] = el_fileGrp 
+                # Assign an empty dictionary that will hold the files of the added fileGrp
+                self._file_cache[fileGrp] = {}
+
         return el_fileGrp
 
     def rename_file_group(self, old, new):
@@ -249,6 +418,11 @@ def rename_file_group(self, old, new):
             raise FileNotFoundError("No such fileGrp '%s'" % old)
         el_fileGrp.set('USE', new)
 
+        # Rename the fileGrp in both caches
+        if self._cache_flag:
+            self._fileGrp_cache[new] = self._fileGrp_cache.pop(old)
+            self._file_cache[new] = deepcopy(self._file_cache.pop(old))
+
     def remove_file_group(self, USE, recursive=False, force=False):
         """
         Remove a ``mets:fileGrp`` (single fixed ``@USE`` or multiple regex ``@USE``)
@@ -285,6 +459,17 @@ def remove_file_group(self, USE, recursive=False, force=False):
                 raise Exception("fileGrp %s is not empty and recursive wasn't set" % USE)
             for f in files:
                 self.remove_one_file(f.get('ID'))
+
+        # Remove the fileGrp from the caches
+        if self._cache_flag:
+            del self._fileGrp_cache[el_fileGrp.get('USE')]
+
+            # Note: Since the files inside the group are removed
+            # with the 'remove_one_file method' above, 
+            # we should not take care of that again.
+            # We just remove the fileGrp.
+            del self._file_cache[el_fileGrp.get('USE')]
+
         el_fileGrp.getparent().remove(el_fileGrp)
 
     def add_file(self, fileGrp, mimetype=None, url=None, ID=None, pageId=None, force=False, local_filename=None, ignore=False, **kwargs):
@@ -310,10 +495,25 @@ def add_file(self, fileGrp, mimetype=None, url=None, ID=None, pageId=None, force
             raise ValueError("Invalid syntax for mets:file/@ID %s (not an xs:ID)" % ID)
         if not REGEX_FILE_ID.fullmatch(fileGrp):
             raise ValueError("Invalid syntax for mets:fileGrp/@USE %s (not an xs:ID)" % fileGrp)
-        el_fileGrp = self._tree.getroot().find(".//mets:fileGrp[@USE='%s']" % (fileGrp), NS)
+
+        el_fileGrp = None
+
+        # If cache is enabled, check there
+        if self._cache_flag: 
+            if fileGrp in self._fileGrp_cache:
+                el_fileGrp = self._fileGrp_cache[fileGrp]
+
+        # cache is not enabled or fileGrp not in the cache
+        if el_fileGrp is None:
+            el_fileGrp = self._tree.getroot().find(".//mets:fileGrp[@USE='%s']" % (fileGrp), NS)
+
+        # the fileGrp is not in the XML tree as well
         if el_fileGrp is None:
             el_fileGrp = self.add_file_group(fileGrp)
-        mets_file = next(self.find_files(ID=ID), None)
+
+        # Since we are sure that fileGrp parameter is set,
+        # we could send that parameter to find_files for direct search
+        mets_file = next(self.find_files(ID=ID, fileGrp=fileGrp), None)
         if mets_file and not ignore:
             if not force:
                 raise Exception("File with ID='%s' already exists" % ID)
@@ -324,7 +524,13 @@ def add_file(self, fileGrp, mimetype=None, url=None, ID=None, pageId=None, force
             mets_file.local_filename = local_filename
         else:
             kwargs = {k: v for k, v in locals().items() if k in ['url', 'ID', 'mimetype', 'pageId', 'local_filename'] and v}
-            mets_file = OcrdFile(ET.SubElement(el_fileGrp, TAG_METS_FILE), mets=self, **kwargs)
+            el_mets_file = ET.SubElement(el_fileGrp, TAG_METS_FILE)
+            mets_file = OcrdFile(el_mets_file, mets=self, **kwargs)
+
+            # Add the file to the cache
+            if self._cache_flag:
+                # print("add_file> Adding to the cache.[%s]" % ID)
+                self._file_cache[fileGrp].update({ID: el_mets_file})
 
         return mets_file
 
@@ -377,6 +583,16 @@ def remove_one_file(self, ID):
                 log.info("Delete empty page %s", page_div)
                 page_div.getparent().remove(page_div)
 
+        # Remove the file from the file cache
+        if self._cache_flag:
+            parent_use = ocrd_file._el.getparent().get('USE')
+            # Note: if the file is in the XML tree,
+            # it should alse be in the file cache.
+            # Anyway, we perform the checks, then remove
+            if parent_use in self._file_cache:
+                if ocrd_file.ID in self._file_cache[parent_use]:
+                    del self._file_cache[parent_use][ocrd_file.ID]
+
         # Delete the file reference
         # pylint: disable=protected-access
         ocrd_file._el.getparent().remove(ocrd_file._el)

ocrd_models/ocrd_models/ocrd_xml_base.py

@@ -16,11 +16,12 @@ class OcrdXmlDocument():
     Base class for XML documents loaded from either content or filename.
     """
 
-    def __init__(self, filename=None, content=None):
+    def __init__(self, filename=None, content=None, cache_flag=False):
         """
         Args:
             filename (string):
             content (string):
+            cache_flag (bool):
         """
         #  print(self, filename, content)
         if filename is None and content is None:
@@ -34,6 +35,9 @@ def __init__(self, filename=None, content=None):
                 raise Exception('File does not exist: %s' % filename)
             self._tree.parse(filename)
 
+        # Cache enabled - True/False
+        self._cache_flag = cache_flag
+
     def to_xml(self, xmllint=False):
         """
         Serialize all properties as pretty-printed XML

requirements_test.txt

@@ -1,6 +1,7 @@
 autopep8
 pytest >= 4.0.0
 generateDS == 2.35.20
+pytest-benchmark >= 3.2.3
 coverage >= 4.5.2
 sphinx
 sphinx_click