Added helper function zeropad

Author: dhondta

docs/pages/helpers.md

@@ -110,6 +110,7 @@ It also provides some other utility functions, decorators, context managers and
 `ts.xor` | repeated XOR function, also allowing to apply an ordinal offset on XORed characters
 `ts.xor_file` | XOR a file with a given key
 `ts.withrepr` | decorator for modifying the representation of a function
+`ts.zeropad` | decorator for zero-padding the result of either an input function or a value, depending on the output value (binary string is padded with "`0`", normal string is padded with "`\x00`", etc).
 
 -----
 

src/tinyscript/VERSION.txt

@@ -1 +1 @@
-1.30.7
+1.30.8

src/tinyscript/helpers/common.py

@@ -11,7 +11,7 @@
 
 
 __all__ = __features__ = ["dateparse", "human_readable_size", "is_admin", "set_exception", "strings",
-                          "strings_from_file", "urlparse", "urlparse_query", "xor", "xor_file", "withrepr"]
+                          "strings_from_file", "urlparse", "urlparse_query", "xor", "xor_file", "withrepr", "zeropad"]
 
 
 def human_readable_size(size, precision=0):
@@ -181,6 +181,42 @@ def xor_file(filename, key, offset=0):
             cursor += l
 
 
+def zeropad(length, default="\x00"):
+    """ Simple decorator to zero-pad the result of the input function regarding its output type.
+    
+    :param length: desired length
+    
+    Examples:
+      zeropad(5)("ok")          => "ok\x00\x00\x00"
+      zeropad(5)("011")         => "01100"
+      zeropad(5)([0,1,1])       =>  [0,1,1,0,0]
+      zeropad(5)(["0","1","1"]) =>  ["0","1","1","0","0"]
+      
+      @zeropad(5)
+      def test(...):
+          ... # e.g. return "ok" => "ok\x00\x00\x00"
+    """
+    from tinyscript.helpers import is_bin, is_function, is_list
+    def _pad(v):
+        l = len(v)
+        if l == 0:
+            return length * ([default] if is_list(v) else default)
+        if l >= length:
+            return v
+        return v + (length - l) * (["\x00", "0"][l > 0 and is_bin(v)] if isinstance(v, str) else \
+                                   [[["\x00", "0"][l > 0 and is_bin(v)]], [0]][l > 0 and isinstance(v[0], int)] \
+                                   if is_list(v) else None)
+    def _zeropad(f):
+        if not is_function(f):
+            return _pad(f)
+        from functools import wraps
+        @wraps(f)
+        def _wrapper(*a, **kw):
+            return _pad(f(*a, **kw))
+        return _wrapper
+    return _zeropad
+
+
 # https://stackoverflow.com/questions/10875442/possible-to-change-a-functions-repr-in-python
 class __reprwrapper(object):
     def __init__(self, repr, func):

src/tinyscript/helpers/path.py

@@ -7,7 +7,6 @@
 import sys
 from mimetypes import guess_type
 from pathlib2 import Path as BasePath
-from pyminizip import compress_multiple, uncompress
 from shutil import copy, copy2, copytree, rmtree
 from tempfile import gettempdir, NamedTemporaryFile as TempFile
 
@@ -506,6 +505,7 @@ def archive(self, path=None, password=None, ask=False, remove=True, **kwargs):
             dst_list.append(str(f.relative_to(self).dirname))
         # Pyminizip changes the current working directory after creation of an archive ; so backup the current
         #  working directory to restore it after compression
+        from pyminizip import compress_multiple, uncompress
         cwd = os.getcwd()
         compress_multiple(src_list, dst_list, str(dst), password or "", 9)
         os.chdir(cwd)
@@ -538,6 +538,7 @@ def load(self, path=None, password=None, ask=False, remove=True, **kwargs):
         dst = Path(path or str(self.dirname.joinpath(self.stem)), create=True)
         # Pyminizip changes the current working directory after extraction of an archive ; so backup the current
         #  working directory to restore it after decompression
+        from pyminizip import uncompress
         cwd = os.getcwd()
         uncompress(str(self), password or "", str(dst), False)
         os.chdir(cwd)

tests/test_helpers_common.py

@@ -62,6 +62,15 @@ def test_func():
         def test_func2():
             pass
         self.assertTrue(repr(test_func2).startswith("<test_function "))
+        self.assertEqual(zeropad(5)("ok"), "ok\x00\x00\x00")
+        self.assertEqual(zeropad(5)("011"), "01100")
+        self.assertEqual(zeropad(5)([0,1,1]), [0,1,1,0,0])
+        self.assertEqual(zeropad(5)(["0","1","1"]), ["0","1","1","0","0"])
+        self.assertEqual(zeropad(5)(lambda: "ok")(), "ok\x00\x00\x00")
+        self.assertEqual(zeropad(5)(lambda: "011")(), "01100")
+        self.assertEqual(zeropad(5)([]), ["\x00","\x00","\x00","\x00","\x00"])
+        self.assertEqual(zeropad(5, default=0)([]), [0,0,0,0,0])
+        self.assertEqual(zeropad(5, default="0")(""), "00000")
 
     def test_lazy_load_functions(self):
         global DUMMY_CONST