Skip to content

Commit

Permalink
Refactor parse_collections
Browse files Browse the repository at this point in the history
  • Loading branch information
@vemel
vemel committed Oct 26, 2024
1 parent bfc20cc commit 6cd402c
Showing 1 changed file with 92 additions and 77 deletions.
169 changes: 92 additions & 77 deletions mypy_boto3_builder/parsers/parse_collections.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
"""

from boto3.resources.base import ServiceResource as Boto3ServiceResource
from boto3.resources.model import Collection as Boto3Collection

from mypy_boto3_builder.parsers.shape_parser import ShapeParser
from mypy_boto3_builder.service_name import ServiceName
Expand All @@ -15,6 +16,93 @@
from mypy_boto3_builder.utils.strings import get_class_prefix


def parse_collection(
collection: Boto3Collection,
service_name: ServiceName,
resource_name: str,
shape_parser: ShapeParser,
) -> Collection:
"""
Parse Collection from boto3 resource.
Arguments:
collection -- Boto3 collection.
service_name -- Service name.
resource_name -- Resource name.
shape_parser -- Shape parser.
"""
if not collection.resource:
raise ValueError(f"Collection {collection.name} does not have resource")
object_class_name = collection.resource.type
collection_record = Collection(
name=f"{resource_name}{get_class_prefix(collection.name)}Collection",
parent_name=resource_name,
attribute_name=collection.name,
service_name=service_name,
type_annotation=InternalImport(collection.name),
object_class_name=object_class_name,
)
self_type = InternalImport(collection_record.name, stringify=True)

collection_record.methods.append(
Method(
name="all",
arguments=(Argument.self(),),
return_type=self_type,
docstring=(
"Get all items from the collection, optionally"
" with a custom page size and item count limit."
),
)
)
filter_method = shape_parser.get_collection_filter_method(
collection_record.name, collection, self_type
)
filter_method.docstring = (
"Get items from the collection, passing keyword arguments along"
" as parameters to the underlying service operation, which are"
" typically used to filter the results."
)
filter_method.type_ignore = True
collection_record.methods.append(filter_method)
batch_methods = shape_parser.get_collection_batch_methods(collection_record.name, collection)
for batch_method in batch_methods:
batch_method.docstring = "Batch method."
collection_record.methods.append(batch_method)

extra_methods = (
Method(
name="limit",
arguments=[Argument.self(), Argument("count", Type.int)],
return_type=self_type,
docstring=f"Return at most this many {object_class_name}s.",
),
Method(
name="page_size",
arguments=[Argument.self(), Argument("count", Type.int)],
return_type=self_type,
docstring=f"Fetch at most this many {object_class_name}s per service request.",
),
Method(
name="pages",
arguments=[Argument.self()],
return_type=TypeSubscript(
Type.Iterator,
[Type.list(InternalImport(name=object_class_name))],
),
docstring=f"A generator which yields pages of {object_class_name}s.",
),
Method(
name="__iter__",
arguments=[Argument.self()],
return_type=TypeSubscript(Type.Iterator, [InternalImport(name=object_class_name)]),
docstring=f"A generator which yields {object_class_name}s.",
),
)
collection_record.methods.extend(extra_methods)
return collection_record


def parse_collections(
service_name: ServiceName,
resource_name: str,
Expand All @@ -25,88 +113,15 @@ def parse_collections(
Extract collections from boto3 resource.
Arguments:
service_name -- Service name.
resource_name -- Resource name.
resource -- boto3 service resource.
Returns:
A list of Collection structures.
shape_parser -- Shape parser.
"""
result: list[Collection] = []
for collection in resource.meta.resource_model.collections:
if not collection.resource:
continue
object_class_name = collection.resource.type
collection_record = Collection(
name=f"{resource_name}{get_class_prefix(collection.name)}Collection",
parent_name=resource_name,
attribute_name=collection.name,
service_name=service_name,
type_annotation=InternalImport(collection.name),
object_class_name=object_class_name,
)
self_type = InternalImport(collection_record.name, stringify=True)

collection_record.methods.append(
Method(
name="all",
arguments=(Argument.self(),),
return_type=self_type,
docstring=(
"Get all items from the collection, optionally"
" with a custom page size and item count limit."
),
)
)
filter_method = shape_parser.get_collection_filter_method(
collection_record.name, collection, self_type
)
filter_method.docstring = (
"Get items from the collection, passing keyword arguments along"
" as parameters to the underlying service operation, which are"
" typically used to filter the results."
)
filter_method.type_ignore = True
collection_record.methods.append(filter_method)
batch_methods = shape_parser.get_collection_batch_methods(
collection_record.name, collection
)
for batch_method in batch_methods:
batch_method.docstring = "Batch method."
collection_record.methods.append(batch_method)
collection_record.methods.append(
Method(
name="limit",
arguments=[Argument.self(), Argument("count", Type.int)],
return_type=self_type,
docstring=f"Return at most this many {object_class_name}s.",
)
)
collection_record.methods.append(
Method(
name="page_size",
arguments=[Argument.self(), Argument("count", Type.int)],
return_type=self_type,
docstring=f"Fetch at most this many {object_class_name}s per service request.",
)
)
collection_record.methods.append(
Method(
name="pages",
arguments=[Argument.self()],
return_type=TypeSubscript(
Type.Iterator,
[Type.list(InternalImport(name=object_class_name))],
),
docstring=f"A generator which yields pages of {object_class_name}s.",
)
)
collection_record.methods.append(
Method(
name="__iter__",
arguments=[Argument.self()],
return_type=TypeSubscript(Type.Iterator, [InternalImport(name=object_class_name)]),
docstring=f"A generator which yields {object_class_name}s.",
)
)

collection_record = parse_collection(collection, service_name, resource_name, shape_parser)
result.append(collection_record)
return result

0 comments on commit 6cd402c

Please sign in to comment.