300 lines
8.4 KiB
Python
300 lines
8.4 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 timeit import default_timer as timer
|
|
from typing import TYPE_CHECKING, Any, BinaryIO, cast
|
|
|
|
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: BinaryIO) -> bytes:
|
|
if file == cast(BinaryIO, sys.stdin):
|
|
# Python 3 < 3.9 does not honor the binary flag,
|
|
# read from the underlying buffer
|
|
if hasattr(file, "buffer"):
|
|
return cast(BinaryIO, file.buffer).read()
|
|
return file.read() # pragma: no cover
|
|
# do not close the file
|
|
try:
|
|
data = file.read()
|
|
finally:
|
|
file.close()
|
|
return data
|
|
|
|
|
|
def writeall(file: BinaryIO, data: bytes) -> None:
|
|
if file == cast(BinaryIO, sys.stdout):
|
|
# Python 3 does not honor the binary flag,
|
|
# write to the underlying buffer
|
|
if hasattr(file, "buffer"):
|
|
file.buffer.write(data)
|
|
else:
|
|
file.write(data) # pragma: no cover
|
|
# do not close the file
|
|
else:
|
|
try:
|
|
file.write(data)
|
|
finally:
|
|
file.close()
|
|
|
|
|
|
def benchmark(duration: float, input: BinaryIO) -> 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: BinaryIO, altchars: bytes | None, output: BinaryIO) -> None:
|
|
data = readall(input)
|
|
data = pybase64.b64encode(data, altchars)
|
|
writeall(output, data)
|
|
|
|
|
|
def decode(input: BinaryIO, altchars: bytes | None, validate: bool, output: BinaryIO) -> 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 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.add_argument(
|
|
"input", type=argparse.FileType("rb"), 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.add_argument(
|
|
"input", type=argparse.FileType("rb"), 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=argparse.FileType("wb"),
|
|
default=sys.stdout,
|
|
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.add_argument(
|
|
"input", type=argparse.FileType("rb"), 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=argparse.FileType("wb"),
|
|
default=sys.stdout,
|
|
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()
|