Skip to content

Commit

Permalink
Added helper function zeropad
Browse files Browse the repository at this point in the history
  • Loading branch information
dhondta committed Mar 22, 2024
1 parent 2aff7d8 commit 91d1975
Show file tree
Hide file tree
Showing 5 changed files with 50 additions and 3 deletions.
1 change: 1 addition & 0 deletions docs/pages/helpers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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).

-----

Expand Down
2 changes: 1 addition & 1 deletion src/tinyscript/VERSION.txt
Original file line number Diff line number Diff line change
@@ -1 +1 @@
1.30.7
1.30.8
38 changes: 37 additions & 1 deletion src/tinyscript/helpers/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
3 changes: 2 additions & 1 deletion src/tinyscript/helpers/path.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
9 changes: 9 additions & 0 deletions tests/test_helpers_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 91d1975

Please sign in to comment.