Improved internal performance

Author: dhondta

src/codext/__common__.py

@@ -57,6 +57,10 @@
 }
 
 __codecs_registry = []
+# caches depending on the state of the registry (cleared whenever it changes)
+__guess_lists = {}
+__names_table = []
+__lookup_cache = {}
 
 MACROS = {}
 PERS_MACROS = {}
@@ -238,15 +242,22 @@ def add(ename, encode=None, decode=None, pattern=None, text=True, add_to_codecs=
     for exc in kwargs.get('extra_exceptions', []):
         _set_exc(exc)  # create additional custom exceptions as builtins
     glob = currentframe().f_back.f_globals
+    # the pattern is compiled once, as it is matched against every encoding name looked up
+    rx = re.compile(pattern) if pattern else None
+    # category of the codec, from the folder of the module adding it
+    cat = glob.get('__file__', os.path.join("custom", "_")).split(os.path.sep)[-2].rstrip("s")
+    if cat not in CODECS_CATEGORIES:
+        CODECS_CATEGORIES.append(cat)
+    category = kwargs.get('category', cat)
     # search function for the new encoding
     @_with_repr(ename)
     def getregentry(encoding):
-        if encoding != ename and not (pattern and re.match(pattern, encoding)):
+        if encoding != ename and not (rx and rx.match(encoding)):
             return
         fenc, fdec, name = encode, decode, encoding
         # prepare CodecInfo input arguments
         if pattern:
-            m, args, i = re.match(pattern, encoding), [], 1
+            m, args, i = rx.match(encoding), [], 1
             try:
                 while True:
                     try:
@@ -329,11 +340,7 @@ class StreamReader(Codec, codecs.StreamReader):
         ci.parameters['add_to_codecs'] = add_to_codecs
         ci.parameters['pattern'] = pattern
         ci.parameters['text'] = text
-        f = glob.get('__file__', os.path.join("custom", "_"))
-        cat = f.split(os.path.sep)[-2].rstrip("s")
-        if cat not in CODECS_CATEGORIES:
-            CODECS_CATEGORIES.append(cat)
-        ci.parameters['category'] = kwargs.get('category', cat)
+        ci.parameters['category'] = category
         ci.parameters['examples'] = kwargs.get('examples', glob.get('__examples__'))
         ci.parameters['guess'] = kwargs.get('guess', glob.get('__guess__', [ename])) or []
         ci.parameters['module'] = kwargs.get('module', glob.get('__name__'))
@@ -349,6 +356,7 @@ class StreamReader(Codec, codecs.StreamReader):
     if kwargs.get('aliases'):
         getregentry.__aliases__ = list(map(lambda n: re.sub(r"[\s\-]", "_", n), kwargs['aliases']))
     getregentry.__pattern__ = pattern
+    getregentry.__category__ = category
     register(getregentry, add_to_codecs)
     return getregentry
 
@@ -623,6 +631,7 @@ def clear():
     """ Clear codext's local registry of search functions. """
     global __codecs_registry, MACROS, PERS_MACROS  # noqa: F824
     __codecs_registry, MACROS, PERS_MACROS = [], {}, {}
+    __registry_changed()
 codecs.clear = clear
 
 
@@ -705,9 +714,11 @@ def list_encodings(*categories):
                 enc.append(ci.name)
     for search_function in CODECS_OVERWRITTEN + __codecs_registry:
         name = search_function.__name__.replace("_", "-")
-        p = search_function.__pattern__
-        ci = search_function(name) if p is None else search_function(generate_string_from_regex(p))
-        c = "other" if ci is None else ci.parameters['category']
+        c = getattr(search_function, "__category__", None)
+        if c is None:
+            p = getattr(search_function, "__pattern__", None)
+            ci = search_function(name) if p is None else search_function(_sample_name(search_function))
+            c = "other" if ci is None else ci.parameters['category']
         if (len(categories) == 0 or c in categories) and c not in exclude:
             enc.append(name)
     for category in categories:
@@ -731,6 +742,7 @@ def remove(name):
             tbr.append(search_function)
     for search_function in tbr:
         __codecs_registry.remove(search_function)
+    __registry_changed()
     try:
         del MACROS[name]
     except KeyError:
@@ -769,6 +781,7 @@ def reset():
     # restore codext's registry
     else:
         __codecs_registry = CODECS_REGISTRY[:]
+    __registry_changed()
     # restore codext's embedded set of macros
     with open(os.path.join(os.path.dirname(__file__), "macros.json")) as f:
         MACROS = json.load(f)
@@ -897,12 +910,15 @@ def decode(obj, encoding='utf-8', errors='strict'):
 codecs.decode = decode
 
 
+_REPEAT = re.compile(r"\[(\d+)\]$")
+
+
 def encode(obj, encoding='utf-8', errors='strict'):
     """ Custom encode function relying on the hooked lookup function. """
-    n, m = 1, re.search(r"\[(\d+)\]$", encoding)
+    n, m = 1, _REPEAT.search(encoding)
     if m:
         n = int(m.group(1))
-        encoding = re.sub(r"\[(\d+)\]$", "", encoding)
+        encoding = encoding[:m.start()]
     ci = lookup(encoding)
     for i in range(n):
         try:
@@ -919,6 +935,26 @@ def encode(obj, encoding='utf-8', errors='strict'):
 
 def lookup(encoding, macro=True):
     """ Hooked lookup function for searching first for codecs in the local registry of this module. """
+    # the results are cached per encoding name, as searching through the registry is costly ; a macro is only returned
+    #  when allowed, and failures are not cached (a codec or a macro may be added afterwards)
+    try:
+        ci = __lookup_cache[encoding]
+        if macro or not isinstance(ci, CodecMacro):
+            return ci
+    except (KeyError, TypeError):
+        pass
+    ci = __lookup(encoding, macro)
+    try:
+        __lookup_cache[encoding] = ci
+    except TypeError:
+        pass
+    return ci
+codecs.lookup = lookup
+
+
+def __lookup(encoding, macro=True):
+    """ Search for the given encoding, first in the local registry of this module, then with the original lookup
+         function and finally among the macros. """
     # first, try to match the given encoding with codecs' search functions
     for search_function in CODECS_OVERWRITTEN + __codecs_registry:
         codecinfo = search_function(encoding)
@@ -928,7 +964,7 @@ def lookup(encoding, macro=True):
     for search_function in CODECS_OVERWRITTEN + __codecs_registry:
         if search_function.__name__.replace("_", "-") == encoding or \
            encoding in getattr(search_function, "__aliases__", []):
-            codecinfo = search_function(generate_string_from_regex(search_function.__pattern__))
+            codecinfo = search_function(_sample_name(search_function))
             if codecinfo is not None:
                 return codecinfo
     # finally, get a CodecInfo with the original lookup function and refine it with a dictionary of parameters
@@ -945,7 +981,6 @@ def lookup(encoding, macro=True):
         e = LookupError(f"unknown encoding: {encoding}")
         e.__cause__ = e  # stop exception chaining
         raise e
-codecs.lookup = lookup
 
 
 def register(search_function, add_to_codecs=False):
@@ -964,6 +999,7 @@ def register(search_function, add_to_codecs=False):
         except LookupError:
             l = __codecs_registry
         l.append(search_function)
+        __registry_changed()
     if add_to_codecs:
         __orig_register(search_function)
 
@@ -974,37 +1010,69 @@ def __register(search_function):
 codecs.register = __register
 
 
+def __names():
+    """ Get the table of names used for searching encodings, as tuples (names to be matched, resulting name, search
+         function), built once per state of the registry. """
+    if len(__names_table) == 0:
+        for search_function in CODECS_OVERWRITTEN + __codecs_registry:
+            n = search_function.__name__
+            __names_table.append(((n, n.replace("_", "-")), n.replace("_", "-"), search_function))
+        for alias, name in ALIASES.items():
+            __names_table.append(((alias, name), name, None))
+    return __names_table
+
+
+def __registry_changed():
+    """ Clear the caches depending on the state of the registry. """
+    __guess_lists.clear()
+    __lookup_cache.clear()
+    __names_table.clear()
+
+
+def _sample_name(search_function):
+    """ Get the encoding name generated from the pattern of the given search function (used when looking up a codec by
+         its name) ; it is cached on the search function, as generating strings from a regex is costly. """
+    try:
+        return search_function.__sample__
+    except AttributeError:
+        search_function.__sample__ = s = generate_string_from_regex(getattr(search_function, "__pattern__", None))
+        return s
+
+
+def _sample_rounds(search_function):
+    """ Get 5 sets of encoding names generated from the pattern of the given search function, for the extended search ;
+         they are cached on the search function, as generating strings from a regex is costly. """
+    try:
+        return search_function.__samples__
+    except AttributeError:
+        p = getattr(search_function, "__pattern__", None)
+        search_function.__samples__ = r = [list(generate_strings_from_regex(p)) for _ in range(5)]
+        return r
+
+
 def search(encoding_regex, extended=True):
     """ Function similar to lookup but allows to search for an encoding based on a regex instead. It searches this way
          into the local registry but also tries a simple lookup with the original lookup function. """
-    matches = []
-    for search_function in CODECS_OVERWRITTEN + __codecs_registry:
-        n = search_function.__name__
-        for name in [n, n.replace("_", "-")]:
-            if re.search(encoding_regex, name):
-                matches.append(n.replace("_", "-"))
-                continue
-        if extended:
+    matches, rx = [], re.compile(encoding_regex)
+    for names, name, search_function in __names():
+        if any(rx.search(n) for n in names):
+            matches.append(name)
+        if extended and search_function is not None:
             # in some cases, encoding_regex can match a generated string that uses a particular portion of its
             #  generating pattern ; e.g. we expect encoding_regex="uu_" to find "uu" and "uu_codec" while it can also
             #  find "morse" or "atbash" very rarely because of their dynamic patterns and the limited number of randomly
             #  generated strings
             # so, we can use a qualified majority voting to ensure we do not get a "junk" encoding in the list of
-            #  matches ; executing 5 times the string generation for a given codec but adding the codec to the list of
+            #  matches ; using 5 sets of generated strings for a given codec but adding the codec to the list of
             #  matches only if we get at least 3 matches ensures that we consider up to 2 failures that could be
             #  stochastic, therefore drastically decreasing the probability to get a "junk" encoding in the matches list
             c = 0
-            for i in range(5):
-                for s in generate_strings_from_regex(search_function.__pattern__):
-                    if re.search(encoding_regex, s):
-                        c += 1
+            for strings in _sample_rounds(search_function):
+                if any(rx.search(s) for s in strings):
+                    c += 1
+                    if c >= 3:
+                        matches.append(search_function.__name__)
                         break
-                if c >= 3:
-                    matches.append(n)
-                    break
-    for s, n in ALIASES.items():
-        if re.search(encoding_regex, s) or re.search(encoding_regex, n):
-            matches.append(n)
     return sorted(list(set(matches)), key=_human_keys)
 codecs.search = search
 
@@ -1270,6 +1338,20 @@ def __guess(prev_input, input, stop_func, depth, max_depth, min_depth, encodings
                 stop, show, scoring_heuristic, extended, debug)
 
 
+def __guess_list(encoding):
+    """ Get the list of encoding names to be tried when guessing for the given encoding, cached per state of the
+         registry. """
+    try:
+        return __guess_lists[encoding]
+    except KeyError:
+        try:
+            l = lookup(encoding, False).parameters['guess']
+        except Exception:
+            l = []
+        __guess_lists[encoding] = l
+        return l
+
+
 def __make_encodings_dict(include, exclude):
     """ Process encodings inclusion and exclusion lists, listing categories and developping codecs' lists of possible
          encoding names. It also creates a cache with the CodecInfo objects for improving performance. """
@@ -1281,11 +1363,7 @@ def _develop(d, keep=True):
             for enc in ((list_encodings(*cc) if (len(cc) > 0 or keep) and len(sc) == 0 else []) + sc):
                 g = []
                 for e in (search(enc, False) or [enc]):
-                    try:
-                        ci = lookup(e, False)
-                        g.extend(ci.parameters['guess'])
-                    except:
-                        pass
+                    g.extend(__guess_list(e))
                 if enc in g:  # e.g. "rot-1" => ["rot-1", "rot-2", ...] ; only "rot-1" is to be selected
                     l.append(enc)
                 else:         # e.g. "rot"   => ["rot-1", "rot-2", ...] ; all the "rot-N" shall be selected