Skip to content

Commit

Permalink
b2 ls without arguments now lists all buckets
Browse files Browse the repository at this point in the history
  • Loading branch information
mjurbanski-reef committed Mar 16, 2024
1 parent dc98459 commit 0556003
Show file tree
Hide file tree
Showing 7 changed files with 204 additions and 44 deletions.
34 changes: 27 additions & 7 deletions b2/_internal/_cli/b2args.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
Utility functions for adding b2-specific arguments to an argparse parser.
"""
import argparse
import functools

from b2._internal._cli.arg_parser_types import wrap_with_argument_type_error
from b2._internal._cli.argcompleters import b2uri_file_completer
Expand All @@ -30,21 +31,40 @@ def b2id_or_file_like_b2_uri(value: str) -> B2URIBase:


B2ID_OR_B2_URI_ARG_TYPE = wrap_with_argument_type_error(parse_b2_uri)
B2ID_OR_B2_URI_OR_ALL_BUCKETS_ARG_TYPE = wrap_with_argument_type_error(
functools.partial(parse_b2_uri, allow_all_buckets=True)
)
B2ID_OR_FILE_LIKE_B2_URI_ARG_TYPE = wrap_with_argument_type_error(b2id_or_file_like_b2_uri)


def add_b2id_or_b2_uri_argument(parser: argparse.ArgumentParser, name="B2_URI"):
def add_b2id_or_b2_uri_argument(
parser: argparse.ArgumentParser, name="B2_URI", *, allow_all_buckets: bool = False
):
"""
Add B2 URI (b2:// or b2id://) as an argument to the parser.
B2 URI can point to a bucket optionally with a object name prefix (directory)
or a file-like object.
If allow_all_buckets is True, the argument will accept B2 URI pointing to all buckets.
"""
parser.add_argument(
name,
type=B2ID_OR_B2_URI_ARG_TYPE,
help="B2 URI pointing to a bucket, directory or a file."
"e.g. b2://yourBucket, b2://yourBucket/file.txt, b2://yourBucket/folderName/, or b2id://fileId",
).completer = b2uri_file_completer
if allow_all_buckets:
argument_spec = parser.add_argument(
name,
type=B2ID_OR_B2_URI_OR_ALL_BUCKETS_ARG_TYPE,
default=B2URI(""),
nargs="?",
help="B2 URI pointing to a bucket, directory, file or all buckets. "
"e.g. b2://yourBucket, b2://yourBucket/file.txt, b2://yourBucket/folderName/, b2id://fileId, or b2://",
)
else:
argument_spec = parser.add_argument(
name,
type=B2ID_OR_B2_URI_ARG_TYPE,
help="B2 URI pointing to a bucket, directory or a file. "
"e.g. b2://yourBucket, b2://yourBucket/file.txt, b2://yourBucket/folderName/, or b2id://fileId",
)

argument_spec.completer = b2uri_file_completer


def add_b2id_or_file_like_b2_uri_argument(parser: argparse.ArgumentParser, name="B2_URI"):
Expand Down
50 changes: 37 additions & 13 deletions b2/_internal/_utils/uri.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
)
from b2sdk.v2.exception import B2Error

from b2._internal._utils.python_compat import removeprefix, singledispatchmethod
from b2._internal._utils.python_compat import singledispatchmethod


class B2URIBase:
Expand All @@ -39,15 +39,13 @@ class B2URI(B2URIBase):
there is no such thing as "directory" in B2, but it is possible to mimic it by using object names with non-trailing
slashes.
To make it possible, it is highly discouraged to use trailing slashes in object names.
Please note `path` attribute should exclude prefixing slash, i.e. `path` should be empty string for the root of the bucket.
"""

bucket_name: str
path: str = ""

def __post_init__(self):
path = removeprefix(self.path, "/")
self.__dict__["path"] = path # hack for a custom init in frozen dataclass

def __str__(self) -> str:
return f"b2://{self.bucket_name}/{self.path}"

Expand Down Expand Up @@ -82,30 +80,56 @@ def __str__(self) -> str:
return f"b2id://{self.file_id}"


def parse_uri(uri: str) -> Path | B2URI | B2FileIdURI:
def parse_uri(uri: str, *, allow_all_buckets: bool = False) -> Path | B2URI | B2FileIdURI:
"""
Parse URI.
:param uri: string to parse
:param allow_all_buckets: if True, allow `b2://` without a bucket name to refer to all buckets
:return: B2 URI or Path
:raises ValueError: if the URI is invalid
"""
if not uri:
raise ValueError("URI cannot be empty")
parsed = urllib.parse.urlsplit(uri)
if parsed.scheme == "":
return pathlib.Path(uri)
return _parse_b2_uri(uri, parsed)
return _parse_b2_uri(uri, parsed, allow_all_buckets=allow_all_buckets)


def parse_b2_uri(uri: str, *, allow_all_buckets: bool = False) -> B2URI | B2FileIdURI:
"""
Parse B2 URI.
def parse_b2_uri(uri: str) -> B2URI | B2FileIdURI:
:param uri: string to parse
:param allow_all_buckets: if True, allow `b2://` without a bucket name to refer to all buckets
:return: B2 URI
:raises ValueError: if the URI is invalid
"""
parsed = urllib.parse.urlsplit(uri)
return _parse_b2_uri(uri, parsed)
return _parse_b2_uri(uri, parsed, allow_all_buckets=allow_all_buckets)


def _parse_b2_uri(uri, parsed: urllib.parse.SplitResult) -> B2URI | B2FileIdURI:
def _parse_b2_uri(
uri, parsed: urllib.parse.SplitResult, allow_all_buckets: bool = False
) -> B2URI | B2FileIdURI:
if parsed.scheme in ("b2", "b2id"):
path = urllib.parse.urlunsplit(parsed._replace(scheme="", netloc=""))
if not parsed.netloc:
if allow_all_buckets:
if path:
raise ValueError(
f"Invalid B2 URI: all buckets URI doesn't allow non-empty path, but {path!r} was provided"
)
return B2URI(bucket_name="")
raise ValueError(f"Invalid B2 URI: {uri!r}")
elif parsed.password or parsed.username:
raise ValueError(
"Invalid B2 URI: credentials passed using `user@password:` syntax are not supported in URI"
"Invalid B2 URI: credentials passed using `user@password:` syntax is not supported in URI"
)

if parsed.scheme == "b2":
path = urllib.parse.urlunsplit(parsed._replace(scheme="", netloc=""))
return B2URI(bucket_name=parsed.netloc, path=path)
return B2URI(bucket_name=parsed.netloc, path=path.removeprefix("/"))
elif parsed.scheme == "b2id":
return B2FileIdURI(file_id=parsed.netloc)
else:
Expand Down
12 changes: 10 additions & 2 deletions b2/_internal/b2v3/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,24 +32,32 @@ class Ls(B2URIBucketNFolderNameArgMixin, BaseLs):
{NAME} ls --recursive --withWildcard bucketName "*.[ct]sv"
List all info.txt files from buckets bX, where X is any character:
List all info.txt files from directories `b?`, where `?` is any character:
.. code-block::
{NAME} ls --recursive --withWildcard bucketName "b?/info.txt"
List all pdf files from buckets b0 to b9 (including sub-directories):
List all pdf files from directories b0 to b9 (including sub-directories):
.. code-block::
{NAME} ls --recursive --withWildcard bucketName "b[0-9]/*.pdf"
List all buckets:
.. code-block::
{NAME} ls
Requires capability:
- **listFiles**
- **listBuckets** (if bucket name is not provided)
"""
ALLOW_ALL_BUCKETS = True


B2.register_subcommand(AuthorizeAccount)
Expand Down
62 changes: 47 additions & 15 deletions b2/_internal/console_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -657,23 +657,30 @@ def get_b2_uri_from_arg(self, args: argparse.Namespace) -> B2URIBase:


class B2URIBucketNFolderNameArgMixin:
ALLOW_ALL_BUCKETS: bool = False

@classmethod
def _setup_parser(cls, parser):
parser.add_argument('bucketName').completer = bucket_name_completer
parser.add_argument(
'bucketName',
nargs='?' if cls.ALLOW_ALL_BUCKETS else None,
).completer = bucket_name_completer
parser.add_argument('folderName', nargs='?').completer = file_name_completer
super()._setup_parser(parser)

def get_b2_uri_from_arg(self, args: argparse.Namespace) -> B2URI:
return B2URI(args.bucketName, args.folderName or '')
return B2URI(args.bucketName or '', args.folderName or '')


class B2IDOrB2URIMixin:
ALLOW_ALL_BUCKETS: bool = False

@classmethod
def _setup_parser(cls, parser):
add_b2id_or_b2_uri_argument(parser)
add_b2id_or_b2_uri_argument(parser, allow_all_buckets=cls.ALLOW_ALL_BUCKETS)
super()._setup_parser(parser)

def get_b2_uri_from_arg(self, args: argparse.Namespace) -> B2URI:
def get_b2_uri_from_arg(self, args: argparse.Namespace) -> B2URI | B2FileIdURI:
return args.B2_URI


Expand Down Expand Up @@ -2063,13 +2070,17 @@ def _setup_parser(cls, parser):
super()._setup_parser(parser)

def _run(self, args):
buckets = self.api.list_buckets()
if args.json:
self._print_json(list(buckets))
return self.__class__.run_list_buckets(self, json_=args.json)

@classmethod
def run_list_buckets(cls, command: Command, *, json_: bool) -> int:
buckets = command.api.list_buckets()
if json_:
command._print_json(list(buckets))
return 0

for b in buckets:
self._print('%s %-10s %s' % (b.id_, b.type_, b.name))
command._print(f'{b.id_} {b.type_:<10} {b.name}')
return 0


Expand Down Expand Up @@ -2236,8 +2247,8 @@ def _setup_parser(cls, parser):
)
super()._setup_parser(parser)

def _print_files(self, args):
generator = self._get_ls_generator(args)
def _print_files(self, args, b2_uri: B2URI | None = None):
generator = self._get_ls_generator(args, b2_uri=b2_uri)

for file_version, folder_name in generator:
self._print_file_version(args, file_version, folder_name)
Expand All @@ -2253,9 +2264,9 @@ def _print_file_version(
name = escape_control_chars(name)
self._print(name)

def _get_ls_generator(self, args):
def _get_ls_generator(self, args, b2_uri: B2URI | None = None):
b2_uri = b2_uri or self.get_b2_uri_from_arg(args)
try:
b2_uri = self.get_b2_uri_from_arg(args)
yield from self.api.list_file_versions_by_uri(
b2_uri,
latest_only=not args.versions,
Expand Down Expand Up @@ -2303,9 +2314,21 @@ def _setup_parser(cls, parser):
super()._setup_parser(parser)

def _run(self, args):
b2_uri = self.get_b2_uri_from_arg(args)
if args.long and args.json:
raise CommandError('Cannot use --long and --json options together')

if b2_uri == B2URI(""):
for option_name in ('long', 'recursive', 'replication'):
if getattr(args, option_name, False):
raise CommandError(
f'Cannot use --{option_name} option without specifying a bucket name'
)
return ListBuckets.run_list_buckets(self, json_=args.json)

if args.json:
i = -1
for i, (file_version, _) in enumerate(self._get_ls_generator(args)):
for i, (file_version, _) in enumerate(self._get_ls_generator(args, b2_uri=b2_uri)):
if i:
self._print(',', end='')
else:
Expand Down Expand Up @@ -2380,24 +2403,33 @@ class Ls(B2IDOrB2URIMixin, BaseLs):
{NAME} ls --recursive --withWildcard "b2://bucketName/*.[ct]sv"
List all info.txt files from buckets bX, where X is any character:
List all info.txt files from directories named `b?`, where `?` is any character:
.. code-block::
{NAME} ls --recursive --withWildcard "b2://bucketName/b?/info.txt"
List all pdf files from buckets b0 to b9 (including sub-directories):
List all pdf files from directories b0 to b9 (including sub-directories):
.. code-block::
{NAME} ls --recursive --withWildcard "b2://bucketName/b[0-9]/*.pdf"
List all buckets:
.. code-block::
{NAME} ls
Requires capability:
- **listFiles**
- **listBuckets** (if bucket name is not provided)
"""
ALLOW_ALL_BUCKETS = True


class BaseRm(ThreadsMixin, AbstractLsCommand, metaclass=ABCMeta):
Expand Down
1 change: 1 addition & 0 deletions changelog.d/+ls_buckets.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add support for calling `b2 ls` without arguments to list all buckets.
Loading

0 comments on commit 0556003

Please sign in to comment.