Applied minor improvements to codebase

dhondta · dhondta · commit 49c83a7bebf7 · 2022-01-19T17:41:37.000+01:00
diff --git a/codext/__common__.py b/codext/__common__.py
@@ -7,7 +7,7 @@
 import re
 import sys
 from encodings.aliases import aliases as ALIASES
-from functools import reduce, wraps
+from functools import reduce, update_wrapper, wraps
 from importlib import import_module
 from inspect import currentframe
 from itertools import chain, product
@@ -39,8 +39,11 @@
            "remove", "reset", "s2i", "search", "stopfunc", "BytesIO", "_input", "_stripl", "CodecMacro",
            "DARWIN", "LANG", "LINUX", "MASKS", "PY3", "UNIX", "WINDOWS"]
 CODECS_REGISTRY = None
+CODECS_OVERWRITTEN = []
 CODECS_CATEGORIES = ["native", "custom"]
-LANG = getlocale()[0][:2].lower() if getlocale() else None
+LANG = getlocale()
+if LANG:
+    LANG = (LANG[0] or "")[:2].lower()
 MASKS = {
     'a': printable,
     'b': "".join(chr(i) for i in range(256)),
@@ -142,6 +145,20 @@ def __repr__(self):
         return "<codext.CodecMacro object for encoding %s at %#x>" % (self.name, id(self))
 
 
+# inspired from: https://stackoverflow.com/questions/10875442/possible-to-change-a-functions-repr-in-python
+class Repr(object):
+    def __init__(self, name, func):
+        self.__name = name
+        self.__func = func
+        update_wrapper(self, func)
+    
+    def __call__(self, *args, **kwargs):
+        return self.__func(*args, **kwargs)
+    
+    def __repr__(self):
+        return "<search-function %s at 0x%x>" % (self.__name, id(self))
+
+
 def __stdin_pipe():
     """ Stdin pipe read function. """
     try:
@@ -173,6 +190,12 @@ def _stripl(s, st_lines, st_crlf):
     return s
 
 
+def _with_repr(name):
+    def _wrapper(f):
+        return Repr(name, f)
+    return _wrapper
+
+
 def add(ename, encode=None, decode=None, pattern=None, text=True, add_to_codecs=False, **kwargs):
     """ This adds a new codec to the codecs module setting its encode and/or decode functions, eventually dynamically
          naming the encoding with a pattern and with file handling.
@@ -195,6 +218,7 @@ def add(ename, encode=None, decode=None, pattern=None, text=True, add_to_codecs=
         raise ValueError("At least one en/decoding function must be defined")
     glob = currentframe().f_back.f_globals
     # search function for the new encoding
+    @_with_repr(ename)
     def getregentry(encoding):
         if encoding != ename and not (pattern and re.match(pattern, encoding)):
             return
@@ -304,6 +328,7 @@ class StreamReader(Codec, codecs.StreamReader):
         getregentry.__aliases__ = list(map(lambda n: re.sub(r"[\s\-]", "_", n), kwargs['aliases']))
     getregentry.__pattern__ = pattern
     register(getregentry, add_to_codecs)
+    return getregentry
 
 
 def add_macro(mname, *encodings):
@@ -500,7 +525,7 @@ def __get_value(token, position, case_changed=False):
                             return __get_value(token_inv_case, position, True)
                         return error_func(token, position)
                     if isinstance(result, list):
-                        result = random.choice(result)
+                        result = result[0]
                     return result + lsep
                 
                 # if a separator is defined, rely on it by splitting the input text
@@ -567,7 +592,7 @@ def __get_value(token, position, case_changed=False):
         kwargs['repl_minlen_b'] = max(1, min(map(len, map(b, set(smapdict.values()) - {''}))))
     except:
         pass
-    add(ename, __generic_code(), __generic_code(True), **kwargs)
+    return add(ename, __generic_code(), __generic_code(True), **kwargs)
 codecs.add_map = add_map
 
 
@@ -651,17 +676,15 @@ def list_encodings(*categories):
     if (len(categories) == 0 or "native" in categories) and "native" not in exclude:
         for a in set(ALIASES.values()):
             try:
-                __orig_lookup(a)
+                ci = __orig_lookup(a)
             except LookupError:
                 continue
-            enc.append(a)
-    for search_function in __codecs_registry:
+            if lookup(a) is ci:
+                enc.append(ci.name)
+    for search_function in CODECS_OVERWRITTEN + __codecs_registry:
         name = search_function.__name__.replace("_", "-")
         p = search_function.__pattern__
-        if p is None:
-            ci = search_function(name)
-        else:
-            ci = search_function(generate_string_from_regex(p))
+        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']
         if (len(categories) == 0 or c in categories) and c not in exclude:
             enc.append(name)
@@ -834,8 +857,9 @@ def _handle_error(token, position, output=""):
 __orig_register = _codecs.register
 
 
-def __add(ename, encode=None, decode=None, pattern=None, text=True, add_to_codecs=True):
-    add(ename, encode, decode, pattern, text, add_to_codecs)
+def __add(ename, encode=None, decode=None, pattern=None, text=True, **kwargs):
+    kwargs.pop('add_to_codecs', None)
+    return add(ename, encode, decode, pattern, text, True, **kwargs)
 __add.__doc__ = add.__doc__
 codecs.add = __add
 
@@ -862,19 +886,19 @@ 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. """
     # first, try to match the given encoding with codecs' search functions
-    for search_function in __codecs_registry:
+    for search_function in CODECS_OVERWRITTEN + __codecs_registry:
         codecinfo = search_function(encoding)
         if codecinfo is not None:
             return codecinfo
     # then, if a codec name was given, generate an encoding name from its pattern and get the CodecInfo
-    for search_function in __codecs_registry:
+    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__))
             if codecinfo is not None:
                 return codecinfo
+    # finally, get a CodecInfo with the original lookup function and refine it with a dictionary of parameters
     try:
-        # finally, get a CodecInfo with the original lookup function and refine it with a dictionary of parameters
         ci = __orig_lookup(encoding)
         ci.parameters = {'category': "native", 'module': "codecs", 'name': ALIASES.get(ci.name, ci.name)}
         return ci
@@ -898,14 +922,19 @@ def register(search_function, add_to_codecs=False):
                                  to remove the codec later
     """
     if search_function not in __codecs_registry:
-        __codecs_registry.append(search_function)
+        try:
+            __orig_lookup(search_function.__name__)
+            l = CODECS_OVERWRITTEN
+        except LookupError:
+            l = __codecs_registry
+        l.append(search_function)
     if add_to_codecs:
         __orig_register(search_function)
 
 
-def __register(search_function, add_to_codecs=True):
+def __register(search_function):
     """ Same as register(...), but with add_to_codecs set by default to True. """
-    register(search_function, add_to_codecs)
+    register(search_function, True)
 codecs.register = __register
 
 
diff --git a/codext/__init__.py b/codext/__init__.py
@@ -3,6 +3,7 @@
 
 """
 from __future__ import print_function
+from _codecs import lookup as orig_lookup
 from ast import literal_eval
 from six import binary_type, text_type
 
@@ -26,6 +27,12 @@
 reset()
 
 
+# overwritten native codec
+add("uu", lambda i, e="strict": orig_lookup("uu").encode(b(i), e),
+          lambda i, e="strict": orig_lookup("uu").decode(b(i), e),
+          pattern=r"^uu(?:[-_]encode|codec)?$", add_to_codecs=True, category="native")
+
+
 def __literal_eval(o):
     """ Non-failing ast.literal_eval alias function. """
     try:
diff --git a/tests/test_common.py b/tests/test_common.py
@@ -8,7 +8,7 @@
 import json
 import random
 import sys
-from codext.__common__ import PERS_MACROS, PERS_MACROS_FILE
+from codext.__common__ import CODECS_OVERWRITTEN, PERS_MACROS, PERS_MACROS_FILE
 from six import b, binary_type, text_type
 from unittest import TestCase
 
@@ -56,25 +56,25 @@ def test_add_codec(self):
         self.assertRaises(ValueError, codext.add, "test")
         self.assertRaises(ValueError, codext.add, "test", "BAD")
         self.assertRaises(ValueError, codext.add, "test", lambda: None, "BAD")
-        self.assertIsNone(codext.add("dummy", dummy_encode, dummy_decode))
+        self.assertIsNotNone(codext.add("dummy", dummy_encode, dummy_decode))
         self.assertEqual(codext.encode("test", "dummy"), "test")
         ci = codext.lookup("dummy")
         for k in ["add_to_codecs", "category", "examples", "name", "pattern", "text"]:
             self.assertIn(k, ci.parameters.keys())
-        self.assertIsNone(codext.add("dummy_errored", None, dummy_errored_decode, r"dummy_errored(\d+)$"))
+        self.assertIsNotNone(codext.add("dummy_errored", None, dummy_errored_decode, r"dummy_errored(\d+)$"))
         self.assertRaises(AttributeError, codext.lookup, "dummy_errored1")
 
     def test_add_map_codec(self):
         ENCMAP = [{'a': "A", 'b': "B", 'c': "C"}, {'d': "D", 'e': "E", 'f': "F"}, {'g': "G", 'h': "H", 'i': "I"}]
-        self.assertIsNone(codext.add_map("dummy2", ENCMAP, pattern=r"^dummy2(?:[-_]?(\d))?$"))
+        self.assertIsNotNone(codext.add_map("dummy2", ENCMAP, pattern=r"^dummy2(?:[-_]?(\d))?$"))
         self.assertRaises(ValueError, codext.add_map, "dummy2", "BAD_ENCMAP")
         self.assertEqual(codext.encode("abc", "dummy2"), "ABC")
         self.assertEqual(codext.encode("abc", "dummy2-1"), "ABC")
         self.assertEqual(codext.encode("def", "dummy2-2"), "DEF")
         self.assertEqual(codext.encode("ghi", "dummy2-3"), "GHI")
         self.assertRaises(LookupError, codext.encode, "test", "dummy2-4")
         ENCMAP = {'': {'a': "A", 'b': "B"}, r'bad': {'a': "B", 'b': "A"}}
-        self.assertIsNone(codext.add_map("dummy3", ENCMAP, pattern=r"^dummy3([-_]inverted)?$"))
+        self.assertIsNotNone(codext.add_map("dummy3", ENCMAP, pattern=r"^dummy3([-_]inverted)?$"))
         self.assertRaises(LookupError, codext.encode, "test", "dummy3_inverted")
         self.assertRaises(ValueError, codext.add_map, "dummy2", ENCMAP, ignore_case="BAD")
         self.assertRaises(ValueError, codext.add_map, "dummy2", ENCMAP, intype="BAD")
@@ -98,13 +98,13 @@ def test_list_codecs(self):
         self.assertFalse(codext.is_native("base64"))
     
     def test_remove_codec(self):
-        self.assertIsNone(codext.add("dummy", dummy_encode, dummy_decode))
+        self.assertIsNotNone(codext.add("dummy", dummy_encode, dummy_decode))
         self.assertEqual(codext.encode("test", "dummy"), "test")
         self.assertIsNone(codext.remove("dummy"))
         self.assertRaises(LookupError, codext.encode, "test", "dummy")
         # special case, when adding a new codec also to the native codecs registry, then it won't be possible to remove
         #  it afterwards
-        self.assertIsNone(codecs.add("dummy2", dummy_encode, dummy_decode))
+        self.assertIsNotNone(codecs.add("dummy2", dummy_encode, dummy_decode))
         self.assertEqual(codecs.encode("test", "dummy2"), "test")
         self.assertIsNone(codecs.remove("dummy2"))
         self.assertEqual(codecs.encode("test", "dummy2"), "test")
@@ -122,6 +122,7 @@ def test_reset_codecs(self):
         self.assertIsNone(codext.reset())
         self.assertIsNotNone(codext.encode("test", "morse"))
         self.assertRaises(LookupError, codext.encode, "test", "dummy")
+        self.assertTrue(len(CODECS_OVERWRITTEN) > 0)
     
     def test_search_codecs(self):
         self.assertIsNotNone(codext.search("morse"))
