From cc10c7f2b572e9cf943cd668f37f93dd4e46c337 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Mon, 17 Jan 2022 16:58:11 -0800 Subject: [PATCH] --nl and --tsv and --csv for list-bucket, refs #48 --- s3_credentials/cli.py | 38 +++++++++++++++++-- tests/test_s3_credentials.py | 72 ++++++++++++++++++++++++++++++------ 2 files changed, 95 insertions(+), 15 deletions(-) diff --git a/s3_credentials/cli.py b/s3_credentials/cli.py index 2305205..e7a2f4c 100644 --- a/s3_credentials/cli.py +++ b/s3_credentials/cli.py @@ -3,6 +3,7 @@ import boto3 import botocore import click import configparser +from csv import DictWriter import io import itertools import json @@ -61,6 +62,18 @@ def common_boto3_options(fn): return fn +def common_output_options(fn): + for decorator in reversed( + ( + click.option("--nl", help="Output newline-delimited JSON", is_flag=True), + click.option("--csv", help="Output CSV", is_flag=True), + click.option("--tsv", help="Output TSV", is_flag=True), + ) + ): + fn = decorator(fn) + return fn + + @click.group() @click.version_option() def cli(): @@ -681,8 +694,9 @@ def ensure_s3_role_exists(iam, sts): @cli.command() @click.argument("bucket") @click.option("--prefix", help="List keys starting with this prefix") +@common_output_options @common_boto3_options -def list_bucket(bucket, prefix, **boto_options): +def list_bucket(bucket, prefix, nl, csv, tsv, **boto_options): "List content of bucket" s3 = make_client("s3", **boto_options) paginator = s3.get_paginator("list_objects_v2") @@ -698,8 +712,7 @@ def list_bucket(bucket, prefix, **boto_options): except botocore.exceptions.ClientError as e: raise click.ClickException(e) - for line in stream_indented_json(iterate()): - click.echo(line) + output(iterate(), nl, csv, tsv) @cli.command() @@ -766,6 +779,25 @@ def get_object(bucket, key, output, **boto_options): s3.download_fileobj(bucket, key, fp) +def output(iterator, nl, csv, tsv): + if nl: + for item in iterator: + click.echo(json.dumps(item, default=repr)) + elif csv or tsv: + first = next(iterator, None) + if first is None: + return + headers = first.keys() + writer = DictWriter( + sys.stdout, headers, dialect="excel-tab" if tsv else "excel" + ) + writer.writeheader() + writer.writerows(itertools.chain([first], iterator)) + else: + for line in stream_indented_json(iterator): + click.echo(line) + + def stream_indented_json(iterator, indent=2): # We have to iterate two-at-a-time so we can know if we # should output a trailing comma or if we have reached diff --git a/tests/test_s3_credentials.py b/tests/test_s3_credentials.py index db78c26..6d0580d 100644 --- a/tests/test_s3_credentials.py +++ b/tests/test_s3_credentials.py @@ -737,7 +737,56 @@ def test_policy(options, expected): assert json.loads(result.output) == json.loads(expected) -def test_list_bucket(stub_s3): +@pytest.mark.parametrize( + "options,expected", + ( + ( + [], + ( + "[\n" + " {\n" + ' "Key": "yolo-causeway-1.jpg",\n' + ' "LastModified": "2019-12-26 17:00:22+00:00",\n' + ' "ETag": "\\"87abea888b22089cabe93a0e17cf34a4\\"",\n' + ' "Size": 5923104,\n' + ' "StorageClass": "STANDARD"\n' + " },\n" + " {\n" + ' "Key": "yolo-causeway-2.jpg",\n' + ' "LastModified": "2019-12-26 17:00:22+00:00",\n' + ' "ETag": "\\"87abea888b22089cabe93a0e17cf34a4\\"",\n' + ' "Size": 5923104,\n' + ' "StorageClass": "STANDARD"\n' + " }\n" + "]\n" + ), + ), + ( + ["--nl"], + ( + '{"Key": "yolo-causeway-1.jpg", "LastModified": "2019-12-26 17:00:22+00:00", "ETag": "\\"87abea888b22089cabe93a0e17cf34a4\\"", "Size": 5923104, "StorageClass": "STANDARD"}\n' + '{"Key": "yolo-causeway-2.jpg", "LastModified": "2019-12-26 17:00:22+00:00", "ETag": "\\"87abea888b22089cabe93a0e17cf34a4\\"", "Size": 5923104, "StorageClass": "STANDARD"}\n' + ), + ), + ( + ["--tsv"], + ( + "Key\tLastModified\tETag\tSize\tStorageClass\n" + 'yolo-causeway-1.jpg\t2019-12-26 17:00:22+00:00\t"""87abea888b22089cabe93a0e17cf34a4"""\t5923104\tSTANDARD\n' + 'yolo-causeway-2.jpg\t2019-12-26 17:00:22+00:00\t"""87abea888b22089cabe93a0e17cf34a4"""\t5923104\tSTANDARD\n' + ), + ), + ( + ["--csv"], + ( + "Key,LastModified,ETag,Size,StorageClass\n" + 'yolo-causeway-1.jpg,2019-12-26 17:00:22+00:00,"""87abea888b22089cabe93a0e17cf34a4""",5923104,STANDARD\n' + 'yolo-causeway-2.jpg,2019-12-26 17:00:22+00:00,"""87abea888b22089cabe93a0e17cf34a4""",5923104,STANDARD\n' + ), + ), + ), +) +def test_list_bucket(stub_s3, options, expected): stub_s3.add_response( "list_objects_v2", { @@ -748,20 +797,19 @@ def test_list_bucket(stub_s3): "ETag": '"87abea888b22089cabe93a0e17cf34a4"', "Size": 5923104, "StorageClass": "STANDARD", - } + }, + { + "Key": "yolo-causeway-2.jpg", + "LastModified": "2019-12-26 17:00:22+00:00", + "ETag": '"87abea888b22089cabe93a0e17cf34a4"', + "Size": 5923104, + "StorageClass": "STANDARD", + }, ] }, ) runner = CliRunner() with runner.isolated_filesystem(): - result = runner.invoke(cli, ["list-bucket", "test-bucket"]) + result = runner.invoke(cli, ["list-bucket", "test-bucket"] + options) assert result.exit_code == 0 - assert json.loads(result.output) == [ - { - "Key": "yolo-causeway-1.jpg", - "LastModified": "2019-12-26 17:00:22+00:00", - "ETag": '"87abea888b22089cabe93a0e17cf34a4"', - "Size": 5923104, - "StorageClass": "STANDARD", - } - ] + assert result.output == expected