293 lines
8.3 KiB
Python
293 lines
8.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import base64
|
||
|
|
import sys
|
||
|
|
from base64 import b64decode as b64decodeValidate
|
||
|
|
from base64 import encodebytes as b64encodebytes
|
||
|
|
from collections.abc import Sequence
|
||
|
|
from pathlib import Path
|
||
|
|
from timeit import default_timer as timer
|
||
|
|
from typing import TYPE_CHECKING, Any
|
||
|
|
|
||
|
|
import pybase64
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from pybase64._typing import Decode, Encode, EncodeBytes
|
||
|
|
|
||
|
|
|
||
|
|
def bench_one(
|
||
|
|
duration: float,
|
||
|
|
data: bytes,
|
||
|
|
enc: Encode,
|
||
|
|
dec: Decode,
|
||
|
|
encbytes: EncodeBytes,
|
||
|
|
altchars: bytes | None = None,
|
||
|
|
validate: bool = False,
|
||
|
|
) -> None:
|
||
|
|
duration = duration / 2.0
|
||
|
|
|
||
|
|
if not validate and altchars is None:
|
||
|
|
number = 0
|
||
|
|
time = timer()
|
||
|
|
while True:
|
||
|
|
encodedcontent = encbytes(data)
|
||
|
|
number += 1
|
||
|
|
if timer() - time > duration:
|
||
|
|
break
|
||
|
|
iter = number
|
||
|
|
time = timer()
|
||
|
|
while iter > 0:
|
||
|
|
encodedcontent = encbytes(data)
|
||
|
|
iter -= 1
|
||
|
|
time = timer() - time
|
||
|
|
print(
|
||
|
|
"{:<32s} {:9.3f} MB/s ({:,d} bytes -> {:,d} bytes)".format(
|
||
|
|
encbytes.__module__ + "." + encbytes.__name__ + ":",
|
||
|
|
((number * len(data)) / (1024.0 * 1024.0)) / time,
|
||
|
|
len(data),
|
||
|
|
len(encodedcontent),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
number = 0
|
||
|
|
time = timer()
|
||
|
|
while True:
|
||
|
|
encodedcontent = enc(data, altchars=altchars)
|
||
|
|
number += 1
|
||
|
|
if timer() - time > duration:
|
||
|
|
break
|
||
|
|
iter = number
|
||
|
|
time = timer()
|
||
|
|
while iter > 0:
|
||
|
|
encodedcontent = enc(data, altchars=altchars)
|
||
|
|
iter -= 1
|
||
|
|
time = timer() - time
|
||
|
|
print(
|
||
|
|
"{:<32s} {:9.3f} MB/s ({:,d} bytes -> {:,d} bytes)".format(
|
||
|
|
enc.__module__ + "." + enc.__name__ + ":",
|
||
|
|
((number * len(data)) / (1024.0 * 1024.0)) / time,
|
||
|
|
len(data),
|
||
|
|
len(encodedcontent),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
number = 0
|
||
|
|
time = timer()
|
||
|
|
while True:
|
||
|
|
decodedcontent = dec(encodedcontent, altchars=altchars, validate=validate)
|
||
|
|
number += 1
|
||
|
|
if timer() - time > duration:
|
||
|
|
break
|
||
|
|
iter = number
|
||
|
|
time = timer()
|
||
|
|
while iter > 0:
|
||
|
|
decodedcontent = dec(encodedcontent, altchars=altchars, validate=validate)
|
||
|
|
iter -= 1
|
||
|
|
time = timer() - time
|
||
|
|
print(
|
||
|
|
"{:<32s} {:9.3f} MB/s ({:,d} bytes -> {:,d} bytes)".format(
|
||
|
|
dec.__module__ + "." + dec.__name__ + ":",
|
||
|
|
((number * len(data)) / (1024.0 * 1024.0)) / time,
|
||
|
|
len(encodedcontent),
|
||
|
|
len(data),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
assert decodedcontent == data
|
||
|
|
|
||
|
|
|
||
|
|
def readall(file: str) -> bytes:
|
||
|
|
if file == "-":
|
||
|
|
return sys.stdin.buffer.read()
|
||
|
|
return Path(file).read_bytes()
|
||
|
|
|
||
|
|
|
||
|
|
def writeall(file: str, data: bytes) -> None:
|
||
|
|
if file == "-":
|
||
|
|
sys.stdout.buffer.write(data)
|
||
|
|
else:
|
||
|
|
Path(file).write_bytes(data)
|
||
|
|
|
||
|
|
|
||
|
|
def benchmark(duration: float, input: str) -> None:
|
||
|
|
print(__package__ + " " + pybase64.get_version())
|
||
|
|
data = readall(input)
|
||
|
|
for altchars in [None, b"-_"]:
|
||
|
|
for validate in [False, True]:
|
||
|
|
print(f"bench: altchars={altchars!r:s}, validate={validate!r:s}")
|
||
|
|
bench_one(
|
||
|
|
duration,
|
||
|
|
data,
|
||
|
|
pybase64.b64encode,
|
||
|
|
pybase64.b64decode,
|
||
|
|
pybase64.encodebytes,
|
||
|
|
altchars,
|
||
|
|
validate,
|
||
|
|
)
|
||
|
|
bench_one(
|
||
|
|
duration,
|
||
|
|
data,
|
||
|
|
base64.b64encode,
|
||
|
|
b64decodeValidate,
|
||
|
|
b64encodebytes,
|
||
|
|
altchars,
|
||
|
|
validate,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def encode(input: str, altchars: bytes | None, output: str) -> None:
|
||
|
|
data = readall(input)
|
||
|
|
data = pybase64.b64encode(data, altchars)
|
||
|
|
writeall(output, data)
|
||
|
|
|
||
|
|
|
||
|
|
def decode(input: str, altchars: bytes | None, validate: bool, output: str) -> None:
|
||
|
|
data = readall(input)
|
||
|
|
data = pybase64.b64decode(data, altchars, validate)
|
||
|
|
writeall(output, data)
|
||
|
|
|
||
|
|
|
||
|
|
class LicenseAction(argparse.Action):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
option_strings: Sequence[str],
|
||
|
|
dest: str,
|
||
|
|
license: str | None = None,
|
||
|
|
help: str | None = "show license information and exit",
|
||
|
|
):
|
||
|
|
super().__init__(
|
||
|
|
option_strings=option_strings,
|
||
|
|
dest=dest,
|
||
|
|
default=argparse.SUPPRESS,
|
||
|
|
nargs=0,
|
||
|
|
help=help,
|
||
|
|
)
|
||
|
|
self.license = license
|
||
|
|
|
||
|
|
def __call__(
|
||
|
|
self,
|
||
|
|
parser: argparse.ArgumentParser,
|
||
|
|
namespace: argparse.Namespace, # noqa: ARG002
|
||
|
|
values: str | Sequence[Any] | None, # noqa: ARG002
|
||
|
|
option_string: str | None = None, # noqa: ARG002
|
||
|
|
) -> None:
|
||
|
|
print(self.license)
|
||
|
|
parser.exit()
|
||
|
|
|
||
|
|
|
||
|
|
def check_file(value: str, is_input: bool) -> str:
|
||
|
|
if value == "-":
|
||
|
|
return value
|
||
|
|
path = Path(value)
|
||
|
|
if is_input:
|
||
|
|
return str(path.resolve(strict=True))
|
||
|
|
return str(path.parent.resolve(strict=True) / path.name)
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: Sequence[str] | None = None) -> None:
|
||
|
|
# main parser
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
prog=__package__, description=__package__ + " command-line tool."
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"-V",
|
||
|
|
"--version",
|
||
|
|
action="version",
|
||
|
|
version=__package__ + " " + pybase64.get_version(),
|
||
|
|
)
|
||
|
|
parser.add_argument("--license", action=LicenseAction, license=pybase64.get_license_text())
|
||
|
|
# create sub-parsers
|
||
|
|
subparsers = parser.add_subparsers(help="tool help")
|
||
|
|
# benchmark parser
|
||
|
|
benchmark_parser = subparsers.add_parser("benchmark", help="-h for usage")
|
||
|
|
benchmark_parser.add_argument(
|
||
|
|
"-d",
|
||
|
|
"--duration",
|
||
|
|
metavar="D",
|
||
|
|
dest="duration",
|
||
|
|
type=float,
|
||
|
|
default=1.0,
|
||
|
|
help="expected duration for a single encode or decode test",
|
||
|
|
)
|
||
|
|
benchmark_parser.register("type", "input file", lambda s: check_file(s, True))
|
||
|
|
benchmark_parser.add_argument(
|
||
|
|
"input", type="input file", help="input file used for the benchmark"
|
||
|
|
)
|
||
|
|
benchmark_parser.set_defaults(func=benchmark)
|
||
|
|
# encode parser
|
||
|
|
encode_parser = subparsers.add_parser("encode", help="-h for usage")
|
||
|
|
encode_parser.register("type", "input file", lambda s: check_file(s, True))
|
||
|
|
encode_parser.register("type", "output file", lambda s: check_file(s, False))
|
||
|
|
encode_parser.add_argument("input", type="input file", help="input file to be encoded")
|
||
|
|
group = encode_parser.add_mutually_exclusive_group()
|
||
|
|
group.add_argument(
|
||
|
|
"-u",
|
||
|
|
"--url",
|
||
|
|
action="store_const",
|
||
|
|
const=b"-_",
|
||
|
|
dest="altchars",
|
||
|
|
help="use URL encoding",
|
||
|
|
)
|
||
|
|
group.add_argument(
|
||
|
|
"-a",
|
||
|
|
"--altchars",
|
||
|
|
dest="altchars",
|
||
|
|
help="use alternative characters for encoding",
|
||
|
|
)
|
||
|
|
encode_parser.add_argument(
|
||
|
|
"-o",
|
||
|
|
"--output",
|
||
|
|
dest="output",
|
||
|
|
type="output file",
|
||
|
|
default="-",
|
||
|
|
help="encoded output file (default to stdout)",
|
||
|
|
)
|
||
|
|
encode_parser.set_defaults(func=encode)
|
||
|
|
# decode parser
|
||
|
|
decode_parser = subparsers.add_parser("decode", help="-h for usage")
|
||
|
|
decode_parser.register("type", "input file", lambda s: check_file(s, True))
|
||
|
|
decode_parser.register("type", "output file", lambda s: check_file(s, False))
|
||
|
|
decode_parser.add_argument("input", type="input file", help="input file to be decoded")
|
||
|
|
group = decode_parser.add_mutually_exclusive_group()
|
||
|
|
group.add_argument(
|
||
|
|
"-u",
|
||
|
|
"--url",
|
||
|
|
action="store_const",
|
||
|
|
const=b"-_",
|
||
|
|
dest="altchars",
|
||
|
|
help="use URL decoding",
|
||
|
|
)
|
||
|
|
group.add_argument(
|
||
|
|
"-a",
|
||
|
|
"--altchars",
|
||
|
|
dest="altchars",
|
||
|
|
help="use alternative characters for decoding",
|
||
|
|
)
|
||
|
|
decode_parser.add_argument(
|
||
|
|
"-o",
|
||
|
|
"--output",
|
||
|
|
dest="output",
|
||
|
|
type="output file",
|
||
|
|
default="-",
|
||
|
|
help="decoded output file (default to stdout)",
|
||
|
|
)
|
||
|
|
decode_parser.add_argument(
|
||
|
|
"--no-validation",
|
||
|
|
dest="validate",
|
||
|
|
action="store_false",
|
||
|
|
help="disable validation of the input data",
|
||
|
|
)
|
||
|
|
decode_parser.set_defaults(func=decode)
|
||
|
|
# ready, parse
|
||
|
|
if argv is None:
|
||
|
|
argv = sys.argv[1:]
|
||
|
|
if len(argv) == 0:
|
||
|
|
argv = ["-h"]
|
||
|
|
args = vars(parser.parse_args(args=argv))
|
||
|
|
func = args.pop("func")
|
||
|
|
func(**args)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|