[scripts/leveldb-dump] support setting separate key and value encodings

This commit is contained in:
Dmytro Meleshko 2021-04-23 21:38:37 +03:00
parent 44ba053b7c
commit 660316fb14

View file

@ -8,36 +8,42 @@ from sys import stdout
import base64 import base64
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument( encoding_names = ["utf8", "base16", "base32", "base64", "base85"]
"--encode", parser.add_argument("--encoding", "-e", choices=encoding_names, default="utf8")
"-e", parser.add_argument("--key-encoding", choices=encoding_names, default=None)
choices=["utf8", "base16", "base32", "base64", "base85"], parser.add_argument("--value-encoding", choices=encoding_names, default=None)
default="utf8",
)
parser.add_argument("db_path", type=Path) parser.add_argument("db_path", type=Path)
cli_args = parser.parse_args() cli_args = parser.parse_args()
def bytes_to_json(b: bytes) -> Union[str, list[int]]: def bytes_to_json(b: bytes, encoding: str) -> Union[str, list[int]]:
if cli_args.encode == "utf8": if encoding == "utf8":
try: try:
return b.decode("utf8") return b.decode("utf8")
except UnicodeDecodeError: except UnicodeDecodeError:
return list(b) return list(b)
elif cli_args.encode == "base16": elif encoding == "base16":
return base64.b16encode(b).decode("ascii") return base64.b16encode(b).decode("ascii")
elif cli_args.encode == "base32": elif encoding == "base32":
return base64.b32encode(b).decode("ascii") return base64.b32encode(b).decode("ascii")
elif cli_args.encode == "base64": elif encoding == "base64":
return base64.b64encode(b).decode("ascii") return base64.b64encode(b).decode("ascii")
elif cli_args.encode == "base85": elif encoding == "base85":
return base64.b85encode(b).decode("ascii") return base64.b85encode(b).decode("ascii")
else: else:
assert False assert False
key_encoding: str = cli_args.key_encoding or cli_args.encoding
value_encoding: str = cli_args.value_encoding or cli_args.encoding
db = plyvel.DB(str(cli_args.db_path), create_if_missing=False) db = plyvel.DB(str(cli_args.db_path), create_if_missing=False)
with db.iterator() as iterator: with db.iterator() as iterator:
for key, value in iterator: for key, value in iterator:
json.dump({"key": bytes_to_json(key), "value": bytes_to_json(value)}, stdout) json.dump(
{
"key": bytes_to_json(key, key_encoding),
"value": bytes_to_json(value, value_encoding),
},
stdout,
)
stdout.write("\n") stdout.write("\n")