def _get_client_import(self, service_name: ServiceName) -> ExternalImport:
client_import = ExternalImport(
source=ImportString.parent() +
ImportString(ServiceModuleName.client.value),
name=Client.get_class_name(service_name),
)
client_import.import_record = InternalImportRecord(
ServiceModuleName.client, client_import.name)
return client_import
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 _make_async_paginators(self) -> None:
for paginator in self.package.paginators:
paginator.bases = [
ExternalImport(ImportString("aiobotocore", "paginate"), "AioPaginator")
]
paginate_method = paginator.get_method("paginate")
assert isinstance(paginate_method.return_type, TypeSubscript)
paginate_method.return_type.parent = Type.AsyncIterator
def _get_type_func_call(self) -> FakeAnnotation:
if not self.func_call:
raise ValueError(f"Value is not a func call: {self.raw}")
if self.func_call["name"] == "datetime":
return ExternalImport(ImportString("datetime"), "datetime")
if self.func_call["name"] == "StreamingBody":
return ExternalImport(ImportString("botocore", "response"),
"StreamingBody")
if self.func_call["name"] == "EventStream":
return ExternalImport(ImportString("botocore", "eventstream"),
"EventStream")
self.logger.warning(f"Unknown function: {self.raw}, fallback to Any")
return Type.Any
def _make_async_client(self) -> None:
self.package.client.bases = [
ExternalImport(ImportString("aiobotocore", "client"), "AioBaseClient")
]
for method in self.package.client.methods:
if method.name in ["exceptions", "get_waiter", "get_paginator", "can_paginate"]:
continue
method.is_async = True
def _get_resource_meta_class(self, service_name: ServiceName) -> ClassRecord:
return ClassRecord(
name=f"{service_name.class_name}ResourceMeta",
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ResourceMeta",
)
],
attributes=[Attribute("client", self._get_client_import(service_name))],
)
class ServiceResource(ClassRecord):
"""
Boto3 ServiceResource.
"""
name: str = "ServiceResource"
alias_name: str = "ServiceResource"
service_name: ServiceName = ServiceNameCatalog.ec2
boto3_service_resource: Optional[Boto3ServiceResource] = None
collections: List[Collection] = field(default_factory=lambda: [])
sub_resources: List[Resource] = field(default_factory=lambda: [])
bases: List[FakeAnnotation] = field(default_factory=lambda: [
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
])
def __hash__(self) -> int:
return hash(self.service_name)
def get_types(self) -> Set[FakeAnnotation]:
types = super().get_types()
for collection in self.collections:
types.update(collection.get_types())
for sub_resource in self.sub_resources:
types.update(sub_resource.get_types())
return types
def get_all_names(self) -> List[str]:
result = [self.name]
for resource in self.sub_resources:
result.append(resource.name)
for collection in self.get_collections():
result.append(collection.name)
return result
def get_collections(self) -> List[Collection]:
collection_names = [i.name for i in self.collections]
result: List[Collection] = []
result.extend(self.collections)
for resource in self.sub_resources:
for collection in resource.collections:
if collection.name in collection_names:
raise ValueError(
f"Conflicting collections: {collection.name}")
collection_names.append(collection.name)
result.append(collection)
return result
def __init__(self, name: str, service_name: ServiceName):
super().__init__(
name=name,
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
],
use_alias=True,
)
self.service_name: ServiceName = service_name
self.collections: list[Collection] = []
def __init__(self, name: str, docstring: str = ""):
super().__init__(
name=name,
docstring=docstring,
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
],
use_alias=True,
)
self.collections: List[Collection] = []
class TestExternalImport:
def setup_method(self) -> None:
self.result = ExternalImport(ImportString("module"), "name")
self.alias = ExternalImport(ImportString("module"), "name", "alias")
def test_init(self) -> None:
assert self.result.name == "name"
assert self.result.alias == ""
assert self.result.source.render() == "module"
assert self.result.import_record.render() == "from module import name"
assert hash(self.result)
def test_render(self) -> None:
assert self.result.render() == "name"
assert self.alias.render() == "alias"
def test_get_import_record(self) -> None:
assert self.result.get_import_record().render(
) == "from module import name"
assert self.alias.get_import_record().render(
) == "from module import name as alias"
def test_copy(self) -> None:
assert self.result.copy().name == "name"
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=ExternalImport(
source=ImportString(self.service_name.module_name,
ServiceModuleName.paginator.value),
name=self.name,
),
)
class Resource(ClassRecord):
"""
Boto3 ServiceResource sub-Resource.
"""
bases: List[FakeAnnotation] = field(default_factory=lambda: [
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
])
collections: List[Collection] = field(default_factory=lambda: [])
def get_types(self) -> Set[FakeAnnotation]:
types = super().get_types()
for collection in self.collections:
types.update(collection.get_types())
return types
def get_client_method(self) -> Method:
"""
Get `get_waiter` method for `Client`.
"""
return Method(
name="get_waiter",
decorators=[Type.overload],
docstring=self.docstring,
arguments=[
Argument("self", None),
Argument("waiter_name",
TypeLiteral(f"{self.name}Name", [self.waiter_name])),
],
return_type=ExternalImport(
source=ImportString.parent() +
ImportString(ServiceModuleName.waiter.value),
name=self.name,
),
)
class Collection(ClassRecord):
"""
Boto3 ServiceResource or Resource collection.
"""
attribute_name: str = "collection"
parent_name: str = "Parent"
type: FakeAnnotation = Type.Any
bases: List[FakeAnnotation] = field(default_factory=lambda: [
ExternalImport(
source=ImportString("boto3", "resources", "collection"),
name="ResourceCollection",
)
])
use_alias: bool = True
def get_types(self) -> Set[FakeAnnotation]:
types = super().get_types()
types.update(self.type.get_types())
return types
def __init__(
self,
name: str,
service_name: ServiceName,
boto3_service_resource: Boto3ServiceResource,
docstring: str = "",
):
super().__init__(
name=name,
docstring=docstring,
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
],
)
self.service_name = service_name
self.boto3_service_resource = boto3_service_resource
self.collections: List[Collection] = []
self.sub_resources: List[Resource] = []
def __init__(
self,
name: str,
attribute_name: str,
parent_name: str,
type_annotation: FakeAnnotation,
docstring: str = "",
):
super().__init__(
name=name,
use_alias=True,
docstring=docstring,
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "collection"),
name="ResourceCollection",
)
],
)
self.attribute_name = attribute_name
self.parent_name = parent_name
self.type_annotation = type_annotation
def get_types(self) -> Set[FakeAnnotation]:
types: Set[FakeAnnotation] = set()
types.add(
ExternalImport(
source="boto3.resources.base",
name="ServiceResource",
alias="Boto3ServiceResource",
))
types.add(TypeAnnotation(ResourceCollection))
for method in self.methods:
types.update(method.get_types())
for attribute in self.attributes:
types.update(attribute.get_types())
for collection in self.collections:
types.update(collection.get_types())
for sub_resource in self.sub_resources:
types.update(sub_resource.get_types())
for type_annotation in types:
if isinstance(type_annotation, InternalImport):
type_annotation.localize(self.service_name)
return types
def __init__(
self,
name: str,
attribute_name: str,
parent_name: str,
service_name: ServiceName,
type_annotation: FakeAnnotation,
object_class_name: str,
):
super().__init__(
name=name,
use_alias=True,
bases=[
ExternalImport(
source=ImportString("boto3", "resources", "collection"),
name="ResourceCollection",
)
],
)
self.service_name = service_name
self.attribute_name = attribute_name
self.parent_name = parent_name
self.type_annotation = type_annotation
self.object_class_name = object_class_name
class ServiceResource(ClassRecord):
"""
Boto3 ServiceResource.
"""
name: str = "ServiceResource"
alias_name: str = "ServiceResource"
service_name: ServiceName = ServiceNameCatalog.ec2
boto3_service_resource: Optional[Boto3ServiceResource] = None
collections: List[Collection] = field(default_factory=lambda: [])
sub_resources: List[Resource] = field(default_factory=lambda: [])
bases: List[FakeAnnotation] = field(default_factory=lambda: [
ExternalImport(
source=ImportString("boto3", "resources", "base"),
name="ServiceResource",
alias="Boto3ServiceResource",
)
])
def __hash__(self) -> int:
return hash(self.service_name)
def get_types(self) -> Set[FakeAnnotation]:
types = super().get_types()
for collection in self.collections:
types.update(collection.get_types())
for sub_resource in self.sub_resources:
types.update(sub_resource.get_types())
return types
def get_all_names(self) -> List[str]:
result = [self.name]
for resource in self.sub_resources:
result.append(resource.name)
for collection in self.get_collections():
result.append(collection.name)
return result
def get_collections(self) -> List[Collection]:
collection_names = [i.name for i in self.collections]
result: List[Collection] = []
result.extend(self.collections)
for resource in self.sub_resources:
for collection in resource.collections:
if collection.name in collection_names:
raise ValueError(
f"Conflicting collections: {collection.name}")
collection_names.append(collection.name)
result.append(collection)
return result
def get_sub_resources(self) -> List[Resource]:
"""
Get sub-resource in safe order.
Returns:
A list of sub resources.
"""
result: List[Resource] = []
all_names: Set[str] = {i.name for i in self.sub_resources}
added_names: Set[str] = set()
sub_resources = list(self.sub_resources)
sub_resources_list: List[Tuple[Resource, List[InternalImport]]] = []
for sub_resource in sub_resources:
internal_imports = sub_resource.get_internal_imports()
sub_resources_list.append((sub_resource, internal_imports))
sub_resources_list.sort(key=lambda x: len(x[1]))
for sub_resource, internal_imports in sub_resources_list:
for internal_import in internal_imports:
if internal_import.name not in all_names:
continue
if internal_import.name in added_names:
continue
internal_import.stringify = True
internal_import.use_alias = False
result.append(sub_resource)
added_names.add(sub_resource.name)
return result
def parse_boto3_stubs_package(
session: Session, service_names: Iterable[ServiceName],
package_data: type[BasePackageData]) -> Boto3StubsPackage:
"""
Parse data for boto3_stubs package.
Arguments:
session -- boto3 session.
service_names -- All available service names.
Returns:
Boto3StubsPackage structure.
"""
result = Boto3StubsPackage(package_data, service_names=service_names)
for service_name in result.service_names:
result.service_packages.append(
parse_fake_service_package(session, service_name, package_data))
init_arguments = [
Argument("region_name", TypeSubscript(Type.Optional, [Type.str]),
Type.Ellipsis),
Argument("api_version", TypeSubscript(Type.Optional, [Type.str]),
Type.Ellipsis),
Argument("use_ssl", TypeSubscript(Type.Optional, [Type.bool]),
Type.Ellipsis),
Argument(
"verify",
TypeSubscript(Type.Union, [Type.bool, Type.str, Type.none]),
Type.Ellipsis,
),
Argument("endpoint_url", TypeSubscript(Type.Optional, [Type.str]),
Type.Ellipsis),
Argument("aws_access_key_id", TypeSubscript(Type.Optional, [Type.str]),
Type.Ellipsis),
Argument("aws_secret_access_key",
TypeSubscript(Type.Optional, [Type.str]), Type.Ellipsis),
Argument("aws_session_token", TypeSubscript(Type.Optional, [Type.str]),
Type.Ellipsis),
Argument("config", TypeSubscript(Type.Optional, [TypeClass(Config)]),
Type.Ellipsis),
]
client_function_decorators: list[TypeAnnotation] = []
if len(result.service_packages) > 1:
client_function_decorators.append(Type.overload)
for service_package in result.service_packages:
package_name = Boto3StubsPackageData.get_service_package_name(
service_package.service_name)
service_argument = Argument(
"service_name",
TypeLiteral(
service_package.service_name.class_name + "Type",
[service_package.service_name.boto3_name],
),
)
client_function = Function(
name="client",
decorators=client_function_decorators,
docstring="",
arguments=[
service_argument,
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(package_name,
ServiceModuleName.client.value),
name=service_package.client.name,
),
body_lines=["..."],
)
result.init_functions.append(client_function)
result.session_class.methods.append(
Method(
name="client",
decorators=client_function_decorators,
docstring="",
arguments=[
Argument("self", None),
service_argument,
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(package_name,
ServiceModuleName.client.value),
name=service_package.client.name,
),
body_lines=["..."],
))
service_resource_packages = [
i for i in result.service_packages if i.service_resource
]
resource_function_decorators: list[TypeAnnotation] = []
if len(service_resource_packages) > 1:
resource_function_decorators.append(Type.overload)
for service_package in service_resource_packages:
assert service_package.service_resource
package_name = Boto3StubsPackageData.get_service_package_name(
service_package.service_name)
service_argument = Argument(
"service_name",
TypeLiteral(
service_package.service_name.class_name + "Type",
[service_package.service_name.boto3_name],
),
)
resource_function = Function(
name="resource",
decorators=resource_function_decorators,
docstring="",
arguments=[
service_argument,
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(package_name,
ServiceModuleName.service_resource.value),
name=service_package.service_resource.name,
),
body_lines=["..."],
)
result.init_functions.append(resource_function)
result.session_class.methods.append(
Method(
name="resource",
decorators=resource_function_decorators,
docstring="",
arguments=[
Argument("self", None),
service_argument,
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(
package_name,
ServiceModuleName.service_resource.value),
name=service_package.service_resource.name,
),
body_lines=["..."],
))
return result
from mypy_boto3_builder.type_annotations.type_class import TypeClass
from mypy_boto3_builder.type_annotations.type_subscript import TypeSubscript
DOCSTRING_TYPE_MAP: dict[str, FakeAnnotation] = {
"bytes":
Type.bytes,
"blob":
Type.bytes,
"boolean":
Type.bool,
"function":
TypeSubscript(Type.Callable, [Type.Ellipsis, Type.Any]),
"method":
TypeSubscript(Type.Callable, [Type.Ellipsis, Type.Any]),
"botocore or boto3 Client":
ExternalImport(source=ImportString("botocore", "client"),
name="BaseClient"),
"datetime":
TypeClass(datetime),
"timestamp":
TypeClass(datetime),
"dict":
Type.DictStrAny,
"structure":
Type.DictStrAny,
"map":
Type.DictStrAny,
"float":
Type.float,
"double":
Type.float,
"int":
"Bucket": {
"copy": {
"CopySource": s3_copy_source_type
}
},
"Object": {
"copy": {
"CopySource": s3_copy_source_type
}
},
},
ServiceNameCatalog.dynamodb: {
"Table": {
"batch_writer": {
"return":
ExternalImport(ImportString("boto3", "dynamodb", "table"),
"BatchWriter")
}
}
},
}
def _get_from_method_map(
method_name: str,
argument_name: str,
method_type_map: MethodTypeMap,
) -> Optional[FakeAnnotation]:
if method_name in method_type_map:
operation_type_map = method_type_map[method_name]
if argument_name in operation_type_map:
return operation_type_map[argument_name]
from mypy_boto3_builder.type_annotations.external_import import ExternalImport
from mypy_boto3_builder.type_annotations.type_subscript import TypeSubstript
TYPE_MAP: Dict[str, FakeAnnotation] = {
"bytes":
TypeAnnotation(bytes),
"blob":
TypeAnnotation(bytes),
"boolean":
TypeAnnotation(bool),
"function":
TypeSubstript(
TypeAnnotation(Callable),
[TypeAnnotation(...), TypeAnnotation(Any)]),
"botocore or boto3 Client":
ExternalImport(source="botocore.client", name="BaseClient"),
"datetime":
TypeAnnotation(datetime),
"timestamp":
TypeAnnotation(datetime),
"dict":
TypeSubstript(
TypeAnnotation(Dict),
[TypeAnnotation(str), TypeAnnotation(Any)]),
"structure":
TypeSubstript(
TypeAnnotation(Dict),
[TypeAnnotation(str), TypeAnnotation(Any)]),
"map":
TypeSubstript(
TypeAnnotation(Dict),
from mypy_boto3_builder.service_name import ServiceNameCatalog
from mypy_boto3_builder.type_annotations.external_import import ExternalImport
from mypy_boto3_builder.type_annotations.fake_annotation import FakeAnnotation
from mypy_boto3_builder.type_annotations.internal_import import AliasInternalImport
from mypy_boto3_builder.type_annotations.type import Type
from mypy_boto3_builder.type_annotations.type_class import TypeClass
from mypy_boto3_builder.type_annotations.type_subscript import TypeSubscript
DOCSTRING_TYPE_MAP: Dict[str, FakeAnnotation] = {
"bytes": Type.bytes,
"blob": Type.bytes,
"boolean": Type.bool,
"function": TypeSubscript(Type.Callable, [Type.Ellipsis, Type.Any]),
"method": TypeSubscript(Type.Callable, [Type.Ellipsis, Type.Any]),
"botocore or boto3 Client": ExternalImport(
source=ImportString("botocore", "client"), name="BaseClient"
),
"datetime": TypeClass(datetime),
"timestamp": TypeClass(datetime),
"dict": Type.DictStrAny,
"structure": Type.DictStrAny,
"map": Type.DictStrAny,
"float": Type.float,
"double": Type.float,
"int": Type.int,
"integer": Type.int,
"long": Type.int,
"a file-like object": TypeSubscript(Type.IO, [Type.Any]),
"seekable file-like object": TypeSubscript(Type.IO, [Type.Any]),
"list": TypeSubscript(Type.List, [Type.Any]),
"L{botocore.paginate.Paginator}": ExternalImport(
def _make_async_waiters(self) -> None:
for waiter in self.package.waiters:
waiter.bases = [ExternalImport(ImportString("aiobotocore", "waiter"), "AIOWaiter")]
for method in waiter.methods:
method.is_async = True
def parse_boto3_stubs_package(
session: Session,
service_names: List[ServiceName]) -> Boto3StubsPackage:
"""
Parse data for boto3_stubs package.
Arguments:
session -- boto3 session.
service_names -- All available service names.
Returns:
Boto3StubsPackage structure.
"""
result = Boto3StubsPackage(service_names=service_names)
for service_name in result.service_names:
result.service_packages.append(
parse_fake_service_package(session, service_name))
init_arguments = [
Argument("region_name", TypeSubscript(Type.Optional, [Type.str]),
Type.none),
Argument("api_version", TypeSubscript(Type.Optional, [Type.str]),
Type.none),
Argument("use_ssl", TypeSubscript(Type.Optional, [Type.bool]),
Type.none),
Argument("verify",
TypeSubscript(Type.Union, [Type.bool, Type.str, Type.none]),
Type.none),
Argument("endpoint_url", TypeSubscript(Type.Optional, [Type.str]),
Type.none),
Argument("aws_access_key_id", TypeSubscript(Type.Optional, [Type.str]),
Type.none),
Argument("aws_secret_access_key",
TypeSubscript(Type.Optional, [Type.str]), Type.none),
Argument("aws_session_token", TypeSubscript(Type.Optional, [Type.str]),
Type.none),
Argument("config", TypeSubscript(Type.Optional, [TypeClass(Config)]),
Type.none),
]
for service_package in result.service_packages:
client_function = Function(
name="client",
decorators=[Type.overload],
docstring="",
arguments=[
Argument("service_name",
TypeLiteral(service_package.service_name.boto3_name)),
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(service_package.service_name.module_name,
ServiceModuleName.client.value),
name=service_package.client.name,
),
body_lines=["..."],
)
result.init_functions.append(client_function)
result.session_class.methods.append(
Method(
name="client",
decorators=[Type.overload],
docstring="",
arguments=[
Argument("self", None),
Argument(
"service_name",
TypeLiteral(service_package.service_name.boto3_name)),
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(
service_package.service_name.module_name,
ServiceModuleName.client.value),
name=service_package.client.name,
),
body_lines=["..."],
))
for service_package in result.service_packages:
if service_package.service_resource:
client_function = Function(
name="resource",
decorators=[Type.overload],
docstring="",
arguments=[
Argument(
"service_name",
TypeLiteral(service_package.service_name.boto3_name)),
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(
service_package.service_name.module_name,
ServiceModuleName.service_resource.value,
),
name=service_package.service_resource.name,
),
body_lines=["..."],
)
result.init_functions.append(client_function)
result.session_class.methods.append(
Method(
name="resource",
decorators=[Type.overload],
docstring="",
arguments=[
Argument("self", None),
Argument(
"service_name",
TypeLiteral(
service_package.service_name.boto3_name)),
*init_arguments,
],
return_type=ExternalImport(
source=ImportString(
service_package.service_name.module_name,
ServiceModuleName.service_resource.value,
),
name=service_package.service_resource.name,
),
body_lines=["..."],
))
return result
class ShapeParser:
"""
Parser for botocore shape files.
Arguments:
session -- Boto3 session.
service_name -- ServiceName.
"""
# Type map for shape types.
SHAPE_TYPE_MAP = {
"integer": Type.int,
"long": Type.int,
"boolean": Type.bool,
"double": Type.float,
"float": Type.float,
"timestamp": ExternalImport(ImportString("datetime"), "datetime"),
"blob": TypeSubscript(Type.Union, [Type.bytes, Type.IO]),
}
# Alias map fixes added by botocore for documentation build.
# https://github.com/boto/botocore/blob/develop/botocore/handlers.py#L773
# https://github.com/boto/botocore/blob/develop/botocore/handlers.py#L1055
ARGUMENT_ALIASES: Dict[str, Dict[str, Dict[str, str]]] = {
ServiceNameCatalog.cloudsearchdomain.boto3_name: {
"Search": {"return": "returnFields"}
},
ServiceNameCatalog.logs.boto3_name: {"CreateExportTask": {"from": "fromTime"}},
ServiceNameCatalog.ec2.boto3_name: {"*": {"Filter": "Filters"}},
ServiceNameCatalog.s3.boto3_name: {
"PutBucketAcl": {"ContentMD5": "None"},
"PutBucketCors": {"ContentMD5": "None"},
"PutBucketLifecycle": {"ContentMD5": "None"},
"PutBucketLogging": {"ContentMD5": "None"},
"PutBucketNotification": {"ContentMD5": "None"},
"PutBucketPolicy": {"ContentMD5": "None"},
"PutBucketReplication": {"ContentMD5": "None"},
"PutBucketRequestPayment": {"ContentMD5": "None"},
"PutBucketTagging": {"ContentMD5": "None"},
"PutBucketVersioning": {"ContentMD5": "None"},
"PutBucketWebsite": {"ContentMD5": "None"},
"PutObjectAcl": {"ContentMD5": "None"},
},
}
def __init__(self, session: Session, service_name: ServiceName):
loader = session._loader # pylint: disable=protected-access
botocore_session: BotocoreSession = session._session # pylint: disable=protected-access
service_data = botocore_session.get_service_data(service_name.boto3_name)
self.service_name = service_name
self.service_model = ServiceModel(service_data, service_name.boto3_name)
self._typed_dict_map: Dict[str, TypeTypedDict] = {}
self._waiters_shape: Shape = {}
try:
self._waiters_shape = loader.load_service_model(
service_name.boto3_name, "waiters-2"
)
except UnknownServiceError:
pass
self._paginators_shape: Shape = {}
try:
self._paginators_shape = loader.load_service_model(
service_name.boto3_name, "paginators-1"
)
except UnknownServiceError:
pass
self._resources_shape: Shape = {}
try:
self._resources_shape = loader.load_service_model(
service_name.boto3_name, "resources-1"
)
except UnknownServiceError:
pass
self.logger = get_logger()
def _get_operation(self, name: str) -> OperationModel:
return self.service_model.operation_model(name)
def _get_operation_names(self) -> List[str]:
return list(
self.service_model.operation_names
) # pylint: disable=not-an-iterable
def _get_paginator(self, name: str) -> Shape:
try:
return self._paginators_shape["pagination"][name]
except KeyError:
raise ShapeParserError(f"Unknown paginator: {name}")
def _get_service_resource(self) -> Shape:
return self._resources_shape["service"]
def _get_resource_shape(self, name: str) -> Shape:
try:
return self._resources_shape["resources"][name]
except KeyError:
raise ShapeParserError(f"Unknown resource: {name}")
def get_paginator_names(self) -> List[str]:
"""
Get available paginator names.
Returns:
A list of paginator names.
"""
result: List[str] = []
for name in self._paginators_shape.get("pagination", []):
result.append(name)
result.sort()
return result
def _get_argument_alias(self, operation_name: str, argument_name: str) -> str:
service_map = self.ARGUMENT_ALIASES.get(self.service_name.boto3_name)
if not service_map:
return argument_name
operation_map: Dict[str, str] = {}
if "*" in service_map:
operation_map = service_map["*"]
if operation_name in service_map:
operation_map = service_map[operation_name]
if not operation_map:
return argument_name
if argument_name not in operation_map:
return argument_name
return operation_map[argument_name]
def _parse_arguments(
self,
class_name: str,
method_name: str,
operation_name: str,
shape: StructureShape,
) -> List[Argument]:
result: List[Argument] = []
required = shape.required_members
for argument_name, argument_shape in shape.members.items():
argument_alias = self._get_argument_alias(operation_name, argument_name)
if argument_alias == "None":
continue
argument_type_stub = get_method_type_stub(
self.service_name, class_name, method_name, argument_name
)
if argument_type_stub is not None:
argument_type = argument_type_stub
else:
argument_type = self._parse_shape(argument_shape)
argument = Argument(argument_alias, argument_type)
if argument_name not in required:
argument.default = Type.none
result.append(argument)
result.sort(key=lambda x: not x.required)
return result
def _parse_return_type(
self, class_name: str, method_name: str, shape: Optional[Shape]
) -> FakeAnnotation:
argument_type_stub = get_method_type_stub(
self.service_name, class_name, method_name, "return"
)
if argument_type_stub is not None:
return argument_type_stub
if shape:
return self._parse_shape(shape)
return Type.none
def get_client_method_map(self) -> Dict[str, Method]:
"""
Get client methods from shape.
Returns:
A map of method name to Method.
"""
result: Dict[str, Method] = {
"can_paginate": Method(
"can_paginate",
[Argument("self", None), Argument("operation_name", Type.str)],
Type.bool,
),
"generate_presigned_url": Method(
"generate_presigned_url",
[
Argument("self", None),
Argument("ClientMethod", Type.str),
Argument("Params", Type.DictStrAny, Type.none),
Argument("ExpiresIn", Type.int, TypeConstant(3600)),
Argument("HttpMethod", Type.str, Type.none),
],
Type.str,
),
}
for operation_name in self._get_operation_names():
operation_model = self._get_operation(operation_name)
arguments: List[Argument] = [Argument("self", None)]
method_name = xform_name(operation_name)
if operation_model.input_shape is not None:
arguments.extend(
self._parse_arguments(
"Client",
method_name,
operation_name,
operation_model.input_shape,
)
)
return_type = self._parse_return_type(
"Client", method_name, operation_model.output_shape
)
method = Method(
name=method_name, arguments=arguments, return_type=return_type
)
result[method.name] = method
return result
@staticmethod
def _parse_shape_string(shape: StringShape) -> FakeAnnotation:
if not shape.enum:
return Type.str
type_literal = TypeLiteral()
for option in shape.enum:
type_literal.add_literal_child(option)
return type_literal
def _parse_shape_map(self, shape: MapShape) -> FakeAnnotation:
type_subscript = TypeSubscript(Type.Dict)
if shape.key:
type_subscript.add_child(self._parse_shape(shape.key))
else:
type_subscript.add_child(Type.str)
if shape.value:
type_subscript.add_child(self._parse_shape(shape.value))
else:
type_subscript.add_child(Type.Any)
return type_subscript
def _parse_shape_structure(self, shape: StructureShape) -> FakeAnnotation:
if not shape.members.items():
return Type.DictStrAny
required = shape.required_members
typed_dict_name = f"{shape.name}TypeDef"
shape_type_stub = get_shape_type_stub(self.service_name, typed_dict_name)
if shape_type_stub:
return shape_type_stub
if typed_dict_name in self._typed_dict_map:
return self._typed_dict_map[typed_dict_name]
typed_dict = TypeTypedDict(typed_dict_name)
self._typed_dict_map[typed_dict_name] = typed_dict
for attr_name, attr_shape in shape.members.items():
typed_dict.add_attribute(
attr_name, self._parse_shape(attr_shape), attr_name in required,
)
return typed_dict
def _parse_shape_list(self, shape: ListShape) -> FakeAnnotation:
type_subscript = TypeSubscript(Type.List)
if shape.member:
type_subscript.add_child(self._parse_shape(shape.member))
else:
type_subscript.add_child(Type.Any)
return type_subscript
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 AliasInternalImport(shape.type_name)
self.logger.warning(f"Unknown shape: {shape}")
return Type.Any
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:
for argument in self._parse_arguments(
"Paginator", "paginate", operation_name, operation_shape.input_shape
):
if argument.name in skip_argument_names:
continue
arguments.append(argument)
arguments.append(Argument("PaginationConfig", paginator_config_type, Type.none))
return_type: FakeAnnotation = Type.none
if operation_shape.output_shape is not None:
return_type = TypeSubscript(
Type.Iterator,
[
self._parse_return_type(
"Paginator", "paginate", operation_shape.output_shape
),
],
)
return Method("paginate", arguments, return_type)
def get_wait_method(self, waiter_name: str) -> Method:
"""
Get Waiter `wait` method.
Arguments:
waiter_name -- Waiter name.
Returns:
Method.
"""
operation_name = self._waiters_shape["waiters"][waiter_name]["operation"]
operation_shape = self._get_operation(operation_name)
arguments: List[Argument] = [Argument("self", None)]
if operation_shape.input_shape is not None:
arguments.extend(
self._parse_arguments(
"Waiter", "wait", operation_name, operation_shape.input_shape
)
)
arguments.append(Argument("WaiterConfig", waiter_config_type, Type.none))
return Method(name="wait", arguments=arguments, return_type=Type.none)
def get_service_resource_method_map(self) -> Dict[str, Method]:
"""
Get methods for ServiceResource.
Returns:
A map of method name to Method.
"""
result: Dict[str, Method] = {
"get_available_subresources": Method(
"get_available_subresources",
[Argument("self", None)],
TypeSubscript(Type.List, [Type.str]),
),
}
service_resource_shape = self._get_service_resource()
for action_name, action_shape in service_resource_shape.get(
"actions", {}
).items():
method = self._get_resource_method(
"ServiceResource", action_name, action_shape
)
result[method.name] = method
return result
def get_resource_method_map(self, resource_name: str) -> Dict[str, Method]:
"""
Get methods for Resource.
Arguments:
resource_name -- Resource name.
Returns:
A map of method name to Method.
"""
resource_shape = self._get_resource_shape(resource_name)
result: Dict[str, Method] = {
"get_available_subresources": Method(
"get_available_subresources",
[Argument("self", None)],
TypeSubscript(Type.List, [Type.str]),
),
"load": Method("load", [Argument("self", None)], Type.none),
"reload": Method("reload", [Argument("self", None)], Type.none),
}
for action_name, action_shape in resource_shape.get("actions", {}).items():
method = self._get_resource_method(resource_name, action_name, action_shape)
result[method.name] = method
for waiter_name in resource_shape.get("waiters", {}):
method = Method(
f"wait_until_{xform_name(waiter_name)}",
[Argument("self", None)],
Type.none,
)
result[method.name] = method
return result
def _get_resource_method(
self, resource_name: str, action_name: str, action_shape: Dict[str, Any]
) -> Method:
return_type: FakeAnnotation = Type.none
method_name = xform_name(action_name)
arguments: List[Argument] = [Argument("self", None)]
if "resource" in action_shape:
return_type = self._parse_return_type(
resource_name, method_name, Shape("resource", action_shape["resource"])
)
path = action_shape["resource"].get("path", "")
if path.endswith("[]"):
return_type = TypeSubscript(Type.List, [return_type])
if "request" in action_shape:
operation_name = action_shape["request"]["operation"]
operation_shape = self._get_operation(operation_name)
skip_argument_names: List[str] = [
i["target"]
for i in action_shape["request"].get("params", {})
if i["source"] == "identifier"
]
if operation_shape.input_shape is not None:
for argument in self._parse_arguments(
resource_name,
method_name,
operation_name,
operation_shape.input_shape,
):
if argument.name not in skip_argument_names:
arguments.append(argument)
if operation_shape.output_shape is not None and return_type is Type.none:
operation_return_type = self._parse_shape(operation_shape.output_shape)
return_type = operation_return_type
return Method(name=method_name, arguments=arguments, return_type=return_type)
def get_collection_filter_method(
self, name: str, collection: Collection, self_type: FakeAnnotation
) -> Method:
"""
Get `filter` classmethod for Resource collection.
Arguments:
name -- Collection record name.
collection -- Boto3 Collection.
class_type -- Collection class type annotation.
Returns:
Filter Method record.
"""
arguments: List[Argument] = [Argument("self", None)]
result = Method("filter", arguments, self_type)
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 get_collection_batch_methods(
self, name: str, collection: Collection
) -> List[Method]:
"""
Get batch operations for Resource collection.
Arguments:
name -- Collection record name.
collection -- Boto3 Collection.
class_type -- Collection self type annotation.
Returns:
List of Method records.
"""
result = []
for batch_action in collection.batch_actions:
method = Method(batch_action.name, [Argument("self", None)], Type.none)
result.append(method)
if batch_action.request:
operation_name = batch_action.request.operation
operation_model = self._get_operation(operation_name)
if operation_model.input_shape is not None:
for argument in self._parse_arguments(
name,
batch_action.name,
operation_name,
operation_model.input_shape,
):
if argument.required:
continue
method.arguments.append(argument)
if operation_model.output_shape is not None:
return_type = self._parse_shape(operation_model.output_shape)
method.return_type = return_type
return result
def setup_method(self) -> None:
self.result = ExternalImport(ImportString("module"), "name")
self.alias = ExternalImport(ImportString("module"), "name", "alias")
