codekingpro/portable-devtools

1from __future__ import annotations2 3import argparse4import base645import sys6from base64 import b64decode as b64decodeValidate7from base64 import encodebytes as b64encodebytes8from collections.abc import Sequence9from pathlib import Path10from timeit import default_timer as timer11from typing import TYPE_CHECKING, Any12 13import pybase6414 15if TYPE_CHECKING:16    from pybase64._typing import Decode, Encode, EncodeBytes17 18 19def bench_one(20    duration: float,21    data: bytes,22    enc: Encode,23    dec: Decode,24    encbytes: EncodeBytes,25    altchars: bytes | None = None,26    validate: bool = False,27) -> None:28    duration = duration / 2.029 30    if not validate and altchars is None:31        number = 032        time = timer()33        while True:34            encodedcontent = encbytes(data)35            number += 136            if timer() - time > duration:37                break38        iter = number39        time = timer()40        while iter > 0:41            encodedcontent = encbytes(data)42            iter -= 143        time = timer() - time44        print(45            "{:<32s} {:9.3f} MB/s ({:,d} bytes -> {:,d} bytes)".format(46                encbytes.__module__ + "." + encbytes.__name__ + ":",47                ((number * len(data)) / (1024.0 * 1024.0)) / time,48                len(data),49                len(encodedcontent),50            )51        )52 53    number = 054    time = timer()55    while True:56        encodedcontent = enc(data, altchars=altchars)57        number += 158        if timer() - time > duration:59            break60    iter = number61    time = timer()62    while iter > 0:63        encodedcontent = enc(data, altchars=altchars)64        iter -= 165    time = timer() - time66    print(67        "{:<32s} {:9.3f} MB/s ({:,d} bytes -> {:,d} bytes)".format(68            enc.__module__ + "." + enc.__name__ + ":",69            ((number * len(data)) / (1024.0 * 1024.0)) / time,70            len(data),71            len(encodedcontent),72        )73    )74 75    number = 076    time = timer()77    while True:78        decodedcontent = dec(encodedcontent, altchars=altchars, validate=validate)79        number += 180        if timer() - time > duration:81            break82    iter = number83    time = timer()84    while iter > 0:85        decodedcontent = dec(encodedcontent, altchars=altchars, validate=validate)86        iter -= 187    time = timer() - time88    print(89        "{:<32s} {:9.3f} MB/s ({:,d} bytes -> {:,d} bytes)".format(90            dec.__module__ + "." + dec.__name__ + ":",91            ((number * len(data)) / (1024.0 * 1024.0)) / time,92            len(encodedcontent),93            len(data),94        )95    )96    assert decodedcontent == data97 98 99def readall(file: str) -> bytes:100    if file == "-":101        return sys.stdin.buffer.read()102    return Path(file).read_bytes()103 104 105def writeall(file: str, data: bytes) -> None:106    if file == "-":107        sys.stdout.buffer.write(data)108    else:109        Path(file).write_bytes(data)110 111 112def benchmark(duration: float, input: str) -> None:113    print(__package__ + " " + pybase64.get_version())114    data = readall(input)115    for altchars in [None, b"-_"]:116        for validate in [False, True]:117            print(f"bench: altchars={altchars!r:s}, validate={validate!r:s}")118            bench_one(119                duration,120                data,121                pybase64.b64encode,122                pybase64.b64decode,123                pybase64.encodebytes,124                altchars,125                validate,126            )127            bench_one(128                duration,129                data,130                base64.b64encode,131                b64decodeValidate,132                b64encodebytes,133                altchars,134                validate,135            )136 137 138def encode(input: str, altchars: bytes | None, output: str) -> None:139    data = readall(input)140    data = pybase64.b64encode(data, altchars)141    writeall(output, data)142 143 144def decode(input: str, altchars: bytes | None, validate: bool, output: str) -> None:145    data = readall(input)146    data = pybase64.b64decode(data, altchars, validate)147    writeall(output, data)148 149 150class LicenseAction(argparse.Action):151    def __init__(152        self,153        option_strings: Sequence[str],154        dest: str,155        license: str | None = None,156        help: str | None = "show license information and exit",157    ):158        super().__init__(159            option_strings=option_strings,160            dest=dest,161            default=argparse.SUPPRESS,162            nargs=0,163            help=help,164        )165        self.license = license166 167    def __call__(168        self,169        parser: argparse.ArgumentParser,170        namespace: argparse.Namespace,  # noqa: ARG002171        values: str | Sequence[Any] | None,  # noqa: ARG002172        option_string: str | None = None,  # noqa: ARG002173    ) -> None:174        print(self.license)175        parser.exit()176 177 178def check_file(value: str, is_input: bool) -> str:179    if value == "-":180        return value181    path = Path(value)182    if is_input:183        return str(path.resolve(strict=True))184    return str(path.parent.resolve(strict=True) / path.name)185 186 187def main(argv: Sequence[str] | None = None) -> None:188    # main parser189    parser = argparse.ArgumentParser(190        prog=__package__, description=__package__ + " command-line tool."191    )192    parser.add_argument(193        "-V",194        "--version",195        action="version",196        version=__package__ + " " + pybase64.get_version(),197    )198    parser.add_argument("--license", action=LicenseAction, license=pybase64.get_license_text())199    # create sub-parsers200    subparsers = parser.add_subparsers(help="tool help")201    # benchmark parser202    benchmark_parser = subparsers.add_parser("benchmark", help="-h for usage")203    benchmark_parser.add_argument(204        "-d",205        "--duration",206        metavar="D",207        dest="duration",208        type=float,209        default=1.0,210        help="expected duration for a single encode or decode test",211    )212    benchmark_parser.register("type", "input file", lambda s: check_file(s, True))213    benchmark_parser.add_argument(214        "input", type="input file", help="input file used for the benchmark"215    )216    benchmark_parser.set_defaults(func=benchmark)217    # encode parser218    encode_parser = subparsers.add_parser("encode", help="-h for usage")219    encode_parser.register("type", "input file", lambda s: check_file(s, True))220    encode_parser.register("type", "output file", lambda s: check_file(s, False))221    encode_parser.add_argument("input", type="input file", help="input file to be encoded")222    group = encode_parser.add_mutually_exclusive_group()223    group.add_argument(224        "-u",225        "--url",226        action="store_const",227        const=b"-_",228        dest="altchars",229        help="use URL encoding",230    )231    group.add_argument(232        "-a",233        "--altchars",234        dest="altchars",235        help="use alternative characters for encoding",236    )237    encode_parser.add_argument(238        "-o",239        "--output",240        dest="output",241        type="output file",242        default="-",243        help="encoded output file (default to stdout)",244    )245    encode_parser.set_defaults(func=encode)246    # decode parser247    decode_parser = subparsers.add_parser("decode", help="-h for usage")248    decode_parser.register("type", "input file", lambda s: check_file(s, True))249    decode_parser.register("type", "output file", lambda s: check_file(s, False))250    decode_parser.add_argument("input", type="input file", help="input file to be decoded")251    group = decode_parser.add_mutually_exclusive_group()252    group.add_argument(253        "-u",254        "--url",255        action="store_const",256        const=b"-_",257        dest="altchars",258        help="use URL decoding",259    )260    group.add_argument(261        "-a",262        "--altchars",263        dest="altchars",264        help="use alternative characters for decoding",265    )266    decode_parser.add_argument(267        "-o",268        "--output",269        dest="output",270        type="output file",271        default="-",272        help="decoded output file (default to stdout)",273    )274    decode_parser.add_argument(275        "--no-validation",276        dest="validate",277        action="store_false",278        help="disable validation of the input data",279    )280    decode_parser.set_defaults(func=decode)281    # ready, parse282    if argv is None:283        argv = sys.argv[1:]284    if len(argv) == 0:285        argv = ["-h"]286    args = vars(parser.parse_args(args=argv))287    func = args.pop("func")288    func(**args)289 290 291if __name__ == "__main__":292    main()293