def add_contextmanager_methods(self) -> None:
"""
Add contextmanager methods.
"""
self.package.client.methods.append(
Method(
"__aenter__",
[
Argument("self", None),
],
return_type=InternalImport(self.package.client.name),
is_async=True,
docstring=self.package.client.docstring,
)
)
self.package.client.methods.append(
Method(
"__aexit__",
[
Argument("self", None),
Argument("exc_type", Type.Any),
Argument("exc_val", Type.Any),
Argument("exc_tb", Type.Any),
],
return_type=Type.Any,
is_async=True,
docstring=self.package.client.docstring,
)
)
def alias_name(self) -> str:
if self._alias_name:
return self._alias_name
if not self.use_alias:
raise ValueError(
f"Cannot get alias for { self.name } with no alias.")
return InternalImport.get_alias(self.name)
def __init__(
self,
name: str,
service_name: ServiceName,
boto3_service_resource: Boto3ServiceResource,
):
self.resource_meta_class = self._get_resource_meta_class(service_name)
super().__init__(
name=name,
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
],
attributes=[
Attribute(
"meta",
InternalImport(
self.resource_meta_class.name,
service_name,
ServiceModuleName.service_resource,
),
)
],
)
self.service_name = service_name
self.boto3_service_resource = boto3_service_resource
self.collections: list[Collection] = []
self.sub_resources: list[Resource] = []
def parse_client(session: Session, service_name: ServiceName, shape_parser: ShapeParser) -> Client:
"""
Parse boto3 client to a structure.
Arguments:
session -- boto3 session.
service_name -- Target service name.
Returns:
Client structure.
"""
client = get_boto3_client(session, service_name)
public_methods = get_public_methods(client)
# remove methods that will be overriden
if "get_paginator" in public_methods:
del public_methods["get_paginator"]
if "get_waiter" in public_methods:
del public_methods["get_waiter"]
result = Client(
name=Client.get_class_name(service_name),
service_name=service_name,
boto3_client=client,
)
shape_method_map = shape_parser.get_client_method_map()
result.methods.append(result.get_exceptions_property())
for method_name, public_method in public_methods.items():
if method_name in shape_method_map:
method = shape_method_map[method_name]
else:
method = parse_method("Client", method_name, public_method, service_name)
docstring = get_short_docstring(inspect.getdoc(public_method) or "")
method.docstring = docstring
result.methods.append(method)
service_model = client.meta.service_model
client_exceptions = ClientExceptionsFactory().create_client_exceptions(service_model)
for exception_class_name in dir(client_exceptions):
if exception_class_name.startswith("_"):
continue
if not exception_class_name[0].isupper():
continue
result.exceptions_class.attributes.append(
Attribute(
exception_class_name,
TypeSubscript(
Type.Type,
[InternalImport("BotocoreClientError", stringify=False)],
),
)
)
result.attributes.append(Attribute("meta", TypeClass(ClientMeta)))
return result
class TestInternalImport:
def setup_method(self) -> None:
self.result = InternalImport("MyClass")
def test_init(self) -> None:
assert self.result.name == "MyClass"
def test_render(self) -> None:
assert self.result.render() == '"MyClass"'
self.result.stringify = False
assert self.result.render() == "MyClass"
self.result.use_alias = True
assert self.result.render() == "_MyClass"
def test_get_import_record(self) -> None:
assert self.result.get_import_record().render() == ""
def test_copy(self) -> None:
assert self.result.copy().name == "MyClass"
def get_client_method(self) -> Method:
return Method(
name="get_paginator",
decorators=[Type.overload],
docstring=self.docstring,
arguments=[
Argument("self", None),
Argument("operation_name", TypeLiteral(self.operation_name)),
],
return_type=InternalImport(
self.name, module_name=ServiceModuleName.paginator),
)
def parse_resource(
name: str,
resource: Boto3ServiceResource,
service_name: ServiceName,
shape_parser: ShapeParser,
) -> Resource:
"""
Parse boto3 sub Resource data.
Arguments:
resource -- Original boto3 resource.
Returns:
Resource structure.
"""
result = Resource(
name=name,
service_name=service_name,
)
shape_method_map = shape_parser.get_resource_method_map(name)
public_methods = get_resource_public_methods(resource.__class__)
for method_name, public_method in public_methods.items():
if method_name in shape_method_map:
method = shape_method_map[method_name]
else:
method = parse_method(name, method_name, public_method,
service_name)
docstring = get_short_docstring(inspect.getdoc(public_method) or "")
method.docstring = "".join((
f"{docstring}\n\n" if docstring else "",
"[Show boto3 documentation]",
f"({service_name.get_boto3_doc_link(name, method_name)})\n",
"[Show boto3-stubs documentation]",
f"({service_name.get_doc_link('service_resource', name, f'{method_name} method')})",
))
result.methods.append(method)
result.attributes.extend(parse_attributes(service_name, name, resource))
result.attributes.extend(parse_identifiers(resource))
result.attributes.extend(parse_references(resource))
collections = parse_collections(name, resource, service_name, shape_parser)
for collection in collections:
result.collections.append(collection)
result.attributes.append(
Attribute(
collection.attribute_name,
InternalImport(collection.name, service_name, stringify=False),
))
return result
def parse_resource(
name: str,
resource: Boto3ServiceResource,
service_name: ServiceName,
shape_parser: ShapeParser,
) -> Resource:
"""
Parse boto3 sub Resource data.
Arguments:
resource -- Original boto3 resource.
Returns:
Resource structure.
"""
result = Resource(
name=name,
service_name=service_name,
)
shape_method_map = shape_parser.get_resource_method_map(name)
public_methods = get_resource_public_methods(resource.__class__)
for method_name, public_method in public_methods.items():
if method_name in shape_method_map:
method = shape_method_map[method_name]
else:
method = parse_method(name, method_name, public_method,
service_name)
docstring = get_short_docstring(inspect.getdoc(public_method) or "")
method.docstring = docstring
result.methods.append(method)
attributes = parse_attributes(service_name, name, resource, shape_parser)
result.attributes.extend(attributes)
identifiers = parse_identifiers(resource)
result.attributes.extend(identifiers)
references = parse_references(resource)
result.attributes.extend(references)
collections = parse_collections(name, resource, service_name, shape_parser)
for collection in collections:
result.collections.append(collection)
result.attributes.append(
Attribute(
collection.attribute_name,
InternalImport(collection.name, service_name, stringify=False),
))
return result
def parse_collections(
resource: Boto3ServiceResource, ) -> Generator[Collection, None, None]:
for collection in resource.meta.resource_model.collections:
yield Collection(
name=collection.name,
docstring=clean_doc(getdoc(collection)),
type=InternalImport(
name=collection.name,
service_name=ServiceName(resource.meta.service_name),
),
methods=list(
parse_methods(
get_instance_public_methods(
getattr(resource, collection.name)))),
)
def parse_resource(
name: str,
resource: Boto3ServiceResource,
service_name: ServiceName,
shape_parser: ShapeParser,
) -> Resource:
"""
Parse boto3 sub Resource data.
Arguments:
resource -- Original boto3 resource.
Returns:
Resource structure.
"""
result = Resource(
name=name,
docstring=(
f"[{name} documentation]"
f"({service_name.doc_link}.ServiceResource.{name})"
),
)
shape_methods_map = shape_parser.get_resource_method_map(name)
public_methods = get_resource_public_methods(resource.__class__)
for method_name, public_method in public_methods.items():
if method_name in shape_methods_map:
method = shape_methods_map[method_name]
else:
method = parse_method(name, method_name, public_method, service_name)
method.docstring = (
f"[{name}.{method_name} documentation]"
f"({service_name.doc_link}.{name}.{method_name})"
)
result.methods.append(method)
result.attributes.extend(parse_attributes(service_name, name, resource))
result.attributes.extend(parse_identifiers(resource))
collections = parse_collections(name, resource, service_name, shape_parser)
for collection in collections:
result.collections.append(collection)
result.attributes.append(
Attribute(
collection.attribute_name, InternalImport(collection.name, service_name)
)
)
return result
def get_paginate_method(self, paginator_name: str) -> Method:
"""
Get Paginator `paginate` method.
Arguments:
paginator_name -- Paginator name.
Returns:
Method.
"""
operation_name = paginator_name
paginator_shape = self._get_paginator(paginator_name)
operation_shape = self._get_operation(operation_name)
skip_argument_names: list[str] = []
input_token = paginator_shape["input_token"]
if isinstance(input_token, list):
skip_argument_names.extend(input_token)
else:
skip_argument_names.append(input_token)
if "limit_key" in paginator_shape:
skip_argument_names.append(paginator_shape["limit_key"])
arguments: list[Argument] = [Argument("self", None)]
if operation_shape.input_shape is not None:
shape_arguments = self._parse_arguments(
"Paginator",
"paginate",
operation_name,
operation_shape.input_shape,
exclude_names=skip_argument_names,
)
shape_arguments.append(
Argument("PaginationConfig", paginator_config_type,
Type.Ellipsis))
arguments.extend(self._get_kw_flags("paginate", shape_arguments))
arguments.extend(shape_arguments)
return_type: FakeAnnotation = Type.none
if operation_shape.output_shape is not None:
page_iterator_import = InternalImport("_PageIterator",
stringify=False)
return_item = self._parse_return_type("Paginator", "paginate",
operation_shape.output_shape)
return_type = TypeSubscript(page_iterator_import, [return_item])
return Method("paginate", arguments, return_type)
def get_exceptions_property(self) -> Method:
"""
Generate Client exceptions property.
"""
return Method(
name="exceptions",
decorators=[TypeClass(property)],
arguments=[
Argument("self", None),
],
return_type=InternalImport(
name=self.exceptions_class.name,
module_name=ServiceModuleName.client,
service_name=self.service_name,
stringify=False,
),
docstring=f"{self.name} exceptions.",
)
def parse_references(resource: Boto3ServiceResource) -> List[Attribute]:
"""
Extract references from boto3 resource.
Arguments:
resource -- boto3 service resource.
Returns:
A list of Attribute structures.
"""
result: List[Attribute] = []
references = resource.meta.resource_model.references
for reference in references:
if not reference.resource:
continue
result.append(
Attribute(reference.name,
type=InternalImport(reference.resource.type)))
return result
def _parse_shape(self, shape: Shape) -> FakeAnnotation:
if shape.type_name in self.SHAPE_TYPE_MAP:
return self.SHAPE_TYPE_MAP[shape.type_name]
if isinstance(shape, StringShape):
return self._parse_shape_string(shape)
if isinstance(shape, MapShape):
return self._parse_shape_map(shape)
if isinstance(shape, StructureShape):
return self._parse_shape_structure(shape)
if isinstance(shape, ListShape):
return self._parse_shape_list(shape)
if shape.type_name in self._resources_shape["resources"]:
return InternalImport(shape.type_name)
self.logger.warning(f"Unknown shape: {shape}")
return Type.Any
def parse_references(resource: Boto3ServiceResource) -> list[Attribute]:
"""
Extract references from boto3 resource.
Arguments:
resource -- boto3 service resource.
Returns:
A list of Attribute structures.
"""
result: list[Attribute] = []
references = resource.meta.resource_model.references
for reference in references:
if not reference.resource:
continue
type_annotation: FakeAnnotation = InternalImport(
reference.resource.type)
if reference.resource.path and "[]" in reference.resource.path:
type_annotation = TypeSubscript(Type.List, [type_annotation])
result.append(
Attribute(reference.name, type_annotation=type_annotation))
return result
def get_collection_filter_method(self, name: str,
collection: Collection) -> Method:
"""
Get `filter` classmethod for Resource collection.
Arguments:
name -- Collection record name.
collection -- Boto3 Collection.
Returns:
Filter Method record.
"""
arguments: List[Argument] = [Argument("cls", None)]
result = Method(
"filter",
arguments,
InternalImport(name=name, service_name=self.service_name),
decorators=[Type.classmethod],
)
if not collection.request:
return result
operation_name = collection.request.operation
operation_model = self._get_operation(operation_name)
if operation_model.input_shape is not None:
for argument in self._parse_arguments(
name,
result.name,
operation_name,
operation_model.input_shape,
):
if argument.required:
continue
arguments.append(argument)
return result
def parse_collections(
parent_name: str,
resource: Boto3ServiceResource,
service_name: ServiceName,
shape_parser: ShapeParser,
) -> List[Collection]:
"""
Extract collections from boto3 resource.
Arguments:
resource -- boto3 service resource.
Returns:
A list of Collection structures.
"""
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"{parent_name}{get_class_prefix(collection.name)}Collection",
parent_name=parent_name,
attribute_name=collection.name,
docstring=(
f"[{parent_name}.{collection.name} documentation]"
f"({service_name.doc_link}.{parent_name}.{collection.name})"),
type=InternalImport(collection.name),
)
self_type = InternalImport(collection_record.name, stringify=True)
collection_record.methods.append(
Method("all", [Argument("self", None)], self_type))
filter_method = shape_parser.get_collection_filter_method(
collection_record.name, collection, self_type)
collection_record.methods.append(filter_method)
batch_methods = shape_parser.get_collection_batch_methods(
collection_record.name, collection)
for batch_method in batch_methods:
collection_record.methods.append(batch_method)
collection_record.methods.append(
Method(
"limit",
[Argument("self", None),
Argument("count", Type.int)],
self_type,
))
collection_record.methods.append(
Method(
"page_size",
[Argument("self", None),
Argument("count", Type.int)],
self_type,
))
collection_record.methods.append(
Method(
"pages",
[Argument("self", None)],
TypeSubscript(
Type.Generator,
[
TypeSubscript(
Type.List,
[InternalImport(name=object_class_name)]),
Type.none,
Type.none,
],
),
))
collection_record.methods.append(
Method(
"__iter__",
[Argument("self", None)],
TypeSubscript(
Type.Iterator,
[InternalImport(name=object_class_name)],
),
))
result.append(collection_record)
return result
"bytes or seekable file-like object":
TypeSubstript(
TypeAnnotation(Union),
[TypeAnnotation(bytes), TypeAnnotation(IO)]),
"str or dict":
TypeSubstript(
TypeAnnotation(Union),
[TypeAnnotation(str), TypeAnnotation(Dict)]),
"list(string)":
TypeSubstript(TypeAnnotation(List), [TypeAnnotation(str)]),
"list of str":
TypeSubstript(TypeAnnotation(List), [TypeAnnotation(str)]),
"None":
TypeAnnotation(None),
":py:class:`ec2.NetworkAcl`":
InternalImport("NetworkAcl", ServiceName.ec2),
":py:class:`EC2.NetworkAcl`":
InternalImport("NetworkAcl", ServiceName.ec2),
"list(:py:class:`ec2.InternetGateway`)":
TypeSubstript(TypeAnnotation(List),
[InternalImport("InternetGateway", ServiceName.ec2)]),
":py:class:`iam.UserPolicy`":
InternalImport("UserPolicy", ServiceName.iam),
":py:class:`IAM.UserPolicy`":
InternalImport("UserPolicy", ServiceName.iam),
"list(:py:class:`iam.VirtualMfaDevice`)":
TypeSubstript(TypeAnnotation(List),
[InternalImport("VirtualMfaDevice", ServiceName.iam)]),
"list(:py:class:`ec2.Image`)":
TypeSubstript(TypeAnnotation(List),
[InternalImport("Image", ServiceName.ec2)]),
def setup_method(self) -> None:
self.result = InternalImport("MyClass")
def parse_collections(
parent_name: str,
resource: Boto3ServiceResource,
service_name: ServiceName,
shape_parser: ShapeParser,
) -> List[Collection]:
"""
Extract collections from boto3 resource.
Arguments:
resource -- boto3 service resource.
Returns:
A list of Collection structures.
"""
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"{parent_name}{get_class_prefix(collection.name)}Collection",
parent_name=parent_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", None)],
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(
"limit",
[Argument("self", None),
Argument("count", Type.int)],
self_type,
docstring=f"Return at most this many {object_class_name}s.",
))
collection_record.methods.append(
Method(
"page_size",
[Argument("self", None),
Argument("count", Type.int)],
self_type,
docstring=
f"Fetch at most this many {object_class_name}s per service request.",
))
collection_record.methods.append(
Method(
"pages",
[Argument("self", None)],
TypeSubscript(
Type.Iterator,
[
TypeSubscript(Type.List,
[InternalImport(name=object_class_name)])
],
),
docstring=
f"A generator which yields pages of {object_class_name}s.",
))
collection_record.methods.append(
Method(
"__iter__",
[Argument("self", None)],
TypeSubscript(Type.Iterator,
[InternalImport(name=object_class_name)]),
docstring=f"A generator which yields {object_class_name}s.",
))
result.append(collection_record)
return result
def parse_service_resource(
session: Session, service_name: ServiceName,
shape_parser: ShapeParser) -> Optional[ServiceResource]:
"""
Parse boto3 ServiceResource data.
Arguments:
session -- boto3 session.
service_name -- Target service name.
Returns:
ServiceResource structure or None if service does not have a resource.
"""
service_resource = get_boto3_resource(session, service_name)
if service_resource is None:
return None
logger = get_logger()
logger.debug("Parsing ServiceResource")
result = ServiceResource(
name=f"{service_name.class_name}ServiceResource",
service_name=service_name,
boto3_service_resource=service_resource,
docstring=(f"[{service_name.class_name}.ServiceResource documentation]"
f"({service_name.doc_link}.ServiceResource)"),
)
public_methods = get_public_methods(service_resource)
shape_method_map = shape_parser.get_service_resource_method_map()
for method_name, public_method in public_methods.items():
if method_name in shape_method_map:
method = shape_method_map[method_name]
else:
method = parse_method("ServiceResource", method_name,
public_method, service_name)
method.docstring = (
f"[ServiceResource.{method_name} documentation]"
f"({service_name.doc_link}.ServiceResource.{method_name})")
result.methods.append(method)
logger.debug("Parsing ServiceResource attributes")
result.attributes.extend(
parse_attributes(service_name, "ServiceResource", service_resource))
result.attributes.extend(parse_identifiers(service_resource))
result.attributes.extend(parse_references(service_resource))
logger.debug("Parsing ServiceResource collections")
collections = parse_collections("ServiceResource", service_resource,
service_name, shape_parser)
for collection in collections:
result.collections.append(collection)
result.attributes.append(
Attribute(
collection.attribute_name,
InternalImport(
collection.name,
service_name,
stringify=False,
),
))
for sub_resource in get_sub_resources(session, service_name,
service_resource):
sub_resource_name = sub_resource.__class__.__name__.split(".", 1)[-1]
logger.debug(f"Parsing {sub_resource_name} sub resource")
result.sub_resources.append(
parse_resource(sub_resource_name, sub_resource, service_name,
shape_parser))
return result
def parse_client(session: Session, service_name: ServiceName,
shape_parser: ShapeParser) -> Client:
"""
Parse boto3 client to a structure.
Arguments:
session -- boto3 session.
service_name -- Target service name.
Returns:
Client structure.
"""
client = get_boto3_client(session, service_name)
public_methods = get_public_methods(client)
# remove methods that will be overriden
if "get_paginator" in public_methods:
del public_methods["get_paginator"]
if "get_waiter" in public_methods:
del public_methods["get_waiter"]
result = Client(
name=f"{service_name.class_name}Client",
service_name=service_name,
boto3_client=client,
docstring=(f"[{service_name.class_name}.Client documentation]"
f"({service_name.doc_link}.Client)"),
)
shape_method_map = shape_parser.get_client_method_map()
for method_name, public_method in public_methods.items():
if method_name in shape_method_map:
method = shape_method_map[method_name]
else:
method = parse_method("Client", method_name, public_method,
service_name)
method.docstring = (f"[Client.{method_name} documentation]"
f"({service_name.doc_link}.Client.{method_name})")
result.methods.append(method)
service_model = client.meta.service_model
client_exceptions = ClientExceptionsFactory().create_client_exceptions(
service_model)
for exception_class_name in dir(client_exceptions):
if exception_class_name.startswith("_"):
continue
if not exception_class_name[0].isupper():
continue
result.exceptions_class.attributes.append(
Attribute(
exception_class_name,
TypeSubscript(
Type.Type,
[TypeClass(ClientError, alias="Boto3ClientError")]),
))
result.attributes.append(
Attribute(
"exceptions",
InternalImport(
name=result.exceptions_class.name,
module_name=ServiceModuleName.client,
service_name=service_name,
stringify=False,
),
))
return result
