class Type:
"""
Predefined FakeAnnotation instances.
"""
Union = TypeAnnotation(Union)
Any = TypeAnnotation(Any)
Dict = TypeAnnotation(Dict)
List = TypeAnnotation(List)
Optional = TypeAnnotation(Optional)
Callable = TypeAnnotation(Callable)
IO = TypeAnnotation(IO)
overload = TypeAnnotation(overload)
classmethod = TypeClass(classmethod)
staticmethod = TypeClass(staticmethod)
none = TypeConstant(None)
str = TypeClass(str)
Set = TypeAnnotation(Set)
bool = TypeClass(bool)
bytes = TypeClass(bytes)
bytearray = TypeClass(bytearray)
int = TypeClass(int)
float = TypeClass(float)
Ellipsis = TypeConstant(...)
Generator = TypeAnnotation(Generator)
Decimal = TypeClass(Decimal)
Type = TypeAnnotation(TypingType)
Iterator = TypeAnnotation(Iterator)
ListAny = TypeSubscript(List, [Any])
DictStrAny = TypeSubscript(Dict, [str, Any])
IOBytes = TypeSubscript(IO, [bytes])
def _get_type_list(self) -> TypeSubscript:
if not self.list_items:
return TypeSubscript(Type.List, [Type.Any])
result = TypeSubscript(Type.List)
for item_index, item in enumerate(self.list_items):
prefix = f"{self.prefix}{item_index if item_index else ''}"
result.add_child(TypeValue(prefix, item).get_type())
return result
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 _get_type_union(self) -> FakeAnnotation:
if not self.union_items:
return Type.Any
result = TypeSubscript(Type.Union)
for item_index, item in enumerate(self.union_items):
prefix = f"{self.prefix}{item_index if item_index else ''}"
result.add_child(TypeValue(prefix, item).get_type())
if all(i is result.children[0] for i in result.children):
return result.children[0]
return result
class TestTypeSubscript:
def setup_method(self) -> None:
self.result = TypeSubscript(Type.Dict, [Type.str, Type.int])
def test_init(self) -> None:
assert self.result.parent == Type.Dict
assert self.result.children == [Type.str, Type.int]
assert hash(self.result)
def test_render(self) -> None:
assert self.result.render() == "Dict[str, int]"
assert TypeSubscript(Type.List).render() == "List"
def test_get_import_record(self) -> None:
assert self.result.get_import_record().render() == "from typing import Dict"
def test_get_types(self) -> None:
assert self.result.get_types() == {Type.Dict, Type.str, Type.int}
def test_add_child(self) -> None:
self.result.add_child(Type.bool)
assert len(self.result.children) == 3
assert self.result.children[-1] == Type.bool
def test_is_dict(self) -> None:
assert self.result.is_dict()
assert not TypeSubscript(Type.List).is_dict()
def test_is_list(self) -> None:
assert not self.result.is_list()
assert TypeSubscript(Type.List).is_list()
def test_copy(self) -> None:
assert self.result.copy().parent == Type.Dict
def __init__(self,
name: str,
service_name: ServiceName,
boto3_client: BaseClient,
docstring: str = ""):
super().__init__(name=name, docstring=docstring)
self.service_name = service_name
self.boto3_client = boto3_client
self.exceptions_class = ClassRecord(name="Exceptions")
self.client_error_class = ClassRecord(
name="BotocoreClientError",
attributes=[
Attribute("MSG_TEMPLATE", Type.str),
],
bases=[TypeClass(BaseException)],
methods=[
Method(
name="__init__",
arguments=[
Argument("self", None),
Argument(
"error_response",
TypeSubscript(Type.Dict, [Type.str, Type.Any])),
Argument("operation_name", Type.str),
],
return_type=Type.none,
body_lines=[
"self.response: Dict[str, Any]",
"self.operation_name: str",
],
),
],
)
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 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
def get_optional(cls, wrapped: FakeAnnotation) -> FakeAnnotation:
"""
Get Optional type annotation.
"""
if (isinstance(wrapped, TypeSubscript)
and isinstance(wrapped.parent, TypeAnnotation)
and wrapped.parent.is_union()):
result = wrapped.copy()
result.add_child(cls.none)
return result
return TypeSubscript(cls.Optional, [wrapped])
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 _get_type_dict(self) -> FakeAnnotation:
if not self.dict_items:
return Type.DictStrAny
first_key = self.dict_items[0]["key"]
if first_key in SYNTAX_TYPE_MAP:
result = TypeSubscript(Type.Dict)
result.add_child(SYNTAX_TYPE_MAP[first_key])
result.add_child(
TypeValue(self.service_name, self.prefix,
self.dict_items[0]["value"]).get_type())
return result
typed_dict_name = f"{self.prefix}TypeDef"
shape_type_stub = get_shape_type_stub(self.service_name,
typed_dict_name)
if shape_type_stub:
return shape_type_stub
typed_dict = TypeTypedDict(typed_dict_name)
for item in self.dict_items:
key_name = self._parse_constant(item["key"])
prefix = f"{self.prefix}{key_name}"
typed_dict.add_attribute(
key_name,
TypeValue(self.service_name, prefix, item["value"]).get_type(),
required=False,
)
return typed_dict
def _parse_shape_list(self,
shape: ListShape,
output_child: bool = False) -> FakeAnnotation:
type_subscript = TypeSubscript(
Type.List) if output_child else TypeSubscript(Type.Sequence)
if shape.member:
type_subscript.add_child(
self.parse_shape(shape.member, output_child=output_child))
else:
type_subscript.add_child(Type.Any)
return type_subscript
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.none))
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:
return_type = TypeSubscript(
Type.Iterator,
[
self._parse_return_type("Paginator", "paginate",
operation_shape.output_shape),
],
)
return Method("paginate", arguments, return_type)
def _get_type_literal(self) -> FakeAnnotation:
if not self.literal_items:
raise ValueError(f"Value is not literal: {self.raw}")
items = [TypeValue(self.prefix, i) for i in self.literal_items]
if all(i.is_literal_item() for i in items):
item_constants = [
self._parse_constant(i.value or "") for i in items
]
return TypeLiteral(*item_constants)
item_types = [i.get_type() for i in items]
if all([i is item_types[0] for i in item_types]):
return item_types[0]
return TypeSubscript(Type.Union, item_types)
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])
operation_shape = None
if "request" in action_shape:
operation_name = action_shape["request"]["operation"]
operation_shape = self._get_operation(operation_name)
skip_argument_names = {
self._get_arg_from_target(i["target"])
for i in action_shape["request"].get("params", {})
if i["source"] == "identifier"
}
if operation_shape.input_shape is not None:
shape_arguments = self._parse_arguments(
resource_name,
method_name,
operation_name,
operation_shape.input_shape,
exclude_names=skip_argument_names,
)
arguments.extend(
self._get_kw_flags(method_name, shape_arguments))
arguments.extend(shape_arguments)
if operation_shape.output_shape is not None and return_type is Type.none:
operation_return_type = self.parse_shape(
operation_shape.output_shape, output=True)
return_type = operation_return_type
method = Method(name=method_name,
arguments=arguments,
return_type=return_type)
if operation_shape and operation_shape.input_shape is not None:
method.request_type_annotation = method.get_request_type_annotation(
self._get_typed_dict_name(operation_shape.input_shape,
postfix=resource_name))
return method
class Type:
"""
Predefined FakeAnnotation instances.
"""
Union = TypeAnnotation("Union")
Any = TypeAnnotation("Any")
Dict = TypeAnnotation("Dict")
Mapping = TypeAnnotation("Mapping")
List = TypeAnnotation("List")
Sequence = TypeAnnotation("Sequence")
Optional = TypeAnnotation("Optional")
Callable = TypeAnnotation("Callable")
IO = TypeAnnotation("IO")
overload = TypeAnnotation("overload")
none = TypeConstant(None)
str = TypeClass(str)
Set = TypeAnnotation("Set")
bool = TypeClass(bool)
bytes = TypeClass(bytes)
bytearray = TypeClass(bytearray)
int = TypeClass(int)
float = TypeClass(float)
Ellipsis = TypeConstant(...)
Decimal = TypeClass(Decimal)
Type = TypeAnnotation("Type")
Iterator = TypeAnnotation("Iterator")
AsyncIterator = TypeAnnotation("AsyncIterator")
datetime = TypeClass(datetime)
ListAny = TypeSubscript(List, [Any])
SequenceAny = TypeSubscript(Sequence, [Any])
MappingStrAny = TypeSubscript(Mapping, [str, Any])
DictStrAny = TypeSubscript(Dict, [str, Any])
DictStrStr = TypeSubscript(Dict, [str, str])
IOBytes = TypeSubscript(IO, [bytes])
RemoveArgument = RemoveArgument()
@classmethod
def get_optional(cls, wrapped: FakeAnnotation) -> FakeAnnotation:
"""
Get Optional type annotation.
"""
if (isinstance(wrapped, TypeSubscript)
and isinstance(wrapped.parent, TypeAnnotation)
and wrapped.parent.is_union()):
result = wrapped.copy()
result.add_child(cls.none)
return result
return TypeSubscript(cls.Optional, [wrapped])
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(
name=batch_action.name,
arguments=[Argument("self", None)],
return_type=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:
shape_arguments = self._parse_arguments(
name,
batch_action.name,
operation_name,
operation_model.input_shape,
optional_only=True,
)
method.arguments.extend(
self._get_kw_flags(batch_action.name, shape_arguments))
method.arguments.extend(shape_arguments)
if operation_model.output_shape is not None:
item_return_type = self.parse_shape(
operation_model.output_shape, output=True)
return_type = TypeSubscript(Type.List, [item_return_type])
method.return_type = return_type
return result
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_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.Sequence, [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(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)
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 = TypeSubscript(Type.Dict, [Type.str, Type.int])
from mypy_boto3_builder.type_annotations.type import Type
from mypy_boto3_builder.type_annotations.type_subscript import TypeSubscript
from mypy_boto3_builder.type_maps.typed_dicts import ec2_tag_type, s3_copy_source_type
__all__ = ("get_method_type_stub", )
ArgumentTypeMap = Dict[str, FakeAnnotation]
MethodTypeMap = Dict[str, ArgumentTypeMap]
ClassTypeMap = Dict[str, MethodTypeMap]
ServiceTypeMap = Dict[ServiceName, ClassTypeMap]
TYPE_MAP: ServiceTypeMap = {
ServiceNameCatalog.ec2: {
"*": {
"create_tags": {
"Resources": TypeSubscript(Type.List, [Type.Any]),
"Tags": TypeSubscript(Type.List, [ec2_tag_type]),
"DryRun": Type.bool,
},
},
"Client": {
"delete_tags": {
"Resources": TypeSubscript(Type.List, [Type.Any]),
"Tags": TypeSubscript(Type.List, [ec2_tag_type]),
"DryRun": Type.bool,
},
},
},
ServiceNameCatalog.s3: {
"Client": {
"copy_object": {
from mypy_boto3_builder.structures.argument import Argument
from mypy_boto3_builder.type_annotations.type import Type
from mypy_boto3_builder.type_annotations.type_subscript import TypeSubscript
from mypy_boto3_builder.type_maps.typed_dicts import ec2_tag_type
__all__ = ("get_method_arguments_stub", )
MethodTypeMap = Dict[str, List[Argument]]
ClassTypeMap = Dict[str, MethodTypeMap]
ServiceTypeMap = Dict[ServiceName, ClassTypeMap]
METHOD_MAP: ServiceTypeMap = {
ServiceNameCatalog.ec2: {
"Instance": {
"delete_tags": [
Argument("Resources", TypeSubscript(Type.List, [Type.Any])),
Argument("Tags", TypeSubscript(Type.List, [ec2_tag_type]),
Type.none),
Argument("DryRun", Type.bool, Type.none),
]
},
},
}
def get_method_arguments_stub(service_name: ServiceName, class_name: str,
method_name: str) -> Optional[List[Argument]]:
"""
Get arguments list for method stub.
Arguments:
class ShapeParser:
"""
Parser for botocore shape files.
Arguments:
session -- Boto3 session.
service_name -- ServiceName.
"""
# Type map for shape types.
SHAPE_TYPE_MAP: Mapping[str, FakeAnnotation] = {
"integer":
Type.int,
"long":
Type.int,
"boolean":
Type.bool,
"double":
Type.float,
"float":
Type.float,
"timestamp":
TypeSubscript(Type.Union, [Type.datetime, Type.str]),
"blob":
TypeSubscript(Type.Union,
[Type.bytes, Type.IOBytes,
TypeClass(StreamingBody)]),
"blob_streaming":
TypeSubscript(Type.Union,
[Type.bytes, Type.IOBytes,
TypeClass(StreamingBody)]),
}
OUTPUT_SHAPE_TYPE_MAP: Mapping[str, FakeAnnotation] = {
"timestamp": Type.datetime,
"blob": Type.bytes,
"blob_streaming": TypeClass(StreamingBody),
}
# 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
botocore_session: BotocoreSession = session._session
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: Mapping[str, Any] | None = None
try:
self._waiters_shape = loader.load_service_model(
service_name.boto3_name, "waiters-2")
except UnknownServiceError:
pass
self._paginators_shape: Mapping[str, Any] | None = None
try:
self._paginators_shape = loader.load_service_model(
service_name.boto3_name, "paginators-1")
except UnknownServiceError:
pass
self._resources_shape: Mapping[str, Any] | None = None
try:
self._resources_shape = loader.load_service_model(
service_name.boto3_name, "resources-1")
except UnknownServiceError:
pass
self.logger = get_logger()
self.response_metadata_typed_dict = TypeTypedDict(
"ResponseMetadataTypeDef",
[
TypedDictAttribute("RequestId", Type.str, True),
TypedDictAttribute("HostId", Type.str, True),
TypedDictAttribute("HTTPStatusCode", Type.int, True),
TypedDictAttribute("HTTPHeaders", Type.DictStrStr, True),
TypedDictAttribute("RetryAttempts", Type.int, True),
],
)
self.proxy_operation_model = OperationModel({}, self.service_model)
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)
def _get_paginator(self, name: str) -> dict[str, Any]:
if not self._paginators_shape:
raise ShapeParserError(f"Unknown paginator: {name}")
try:
return self._paginators_shape["pagination"][name]
except KeyError as e:
raise ShapeParserError(f"Unknown paginator: {name}") from e
def _get_service_resource(self) -> dict[str, Any]:
if not self._resources_shape:
raise ShapeParserError("Resource shape not found")
return self._resources_shape["service"]
def _get_resource_shape(self, name: str) -> dict[str, Any]:
if not self._resources_shape:
raise ShapeParserError("Resource shape not found")
try:
return self._resources_shape["resources"][name]
except KeyError as e:
raise ShapeParserError(f"Unknown resource: {name}") from e
def get_paginator_names(self) -> list[str]:
"""
Get available paginator names.
Returns:
A list of paginator names.
"""
result: list[str] = []
if self._paginators_shape:
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,
exclude_names: Iterable[str] = tuple(),
optional_only: bool = False,
) -> list[Argument]:
result: list[Argument] = []
required = shape.required_members
for argument_name, argument_shape in shape.members.items():
if argument_name in exclude_names:
continue
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 Type.RemoveArgument:
continue
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.Ellipsis
if optional_only and argument.required:
continue
# FIXME: https://github.com/boto/boto3/issues/2813
# if not argument.required and argument.type_annotation:
# argument.type_annotation = Type.get_optional(argument.type_annotation)
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: Shape | None) -> 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, output=True)
return Type.none
@staticmethod
def _get_kw_flags(method_name: str,
arguments: Sequence[Argument]) -> list[Argument]:
if len(arguments) and not method_name[0].isupper():
return [Argument.kwflag()]
return []
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.MappingStrAny, Type.Ellipsis),
Argument("ExpiresIn", Type.int, TypeConstant(3600)),
Argument("HttpMethod", Type.str, Type.Ellipsis),
],
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:
shape_arguments = self._parse_arguments(
"Client",
method_name,
operation_name,
operation_model.input_shape,
)
arguments.extend(
self._get_kw_flags(method_name, shape_arguments))
arguments.extend(shape_arguments)
return_type = self._parse_return_type("Client", method_name,
operation_model.output_shape)
method = Method(name=method_name,
arguments=arguments,
return_type=return_type)
if operation_model.input_shape:
method.request_type_annotation = method.get_request_type_annotation(
self._get_typed_dict_name(operation_model.input_shape,
postfix="Request"))
result[method.name] = method
return result
@staticmethod
def _get_typed_dict_name(shape: Shape, postfix: str = "") -> str:
return f"{shape.name}{postfix}TypeDef"
def _parse_shape_string(self, shape: StringShape) -> FakeAnnotation:
if not shape.enum:
return Type.str
literal_name = f"{shape.name}Type"
literal_type_stub = get_literal_type_stub(self.service_name,
literal_name)
if literal_type_stub:
return literal_type_stub
return TypeLiteral(f"{shape.name}Type",
[option for option in shape.enum])
def _parse_shape_map(
self,
shape: MapShape,
output_child: bool = False,
is_streaming: bool = False,
) -> FakeAnnotation:
type_subscript = TypeSubscript(
Type.Dict) if output_child else TypeSubscript(Type.Mapping)
if shape.key:
type_subscript.add_child(
self.parse_shape(shape.key,
output_child=output_child,
is_streaming=is_streaming))
else:
type_subscript.add_child(Type.str)
if shape.value:
type_subscript.add_child(
self.parse_shape(shape.value,
output_child=output_child,
is_streaming=is_streaming))
else:
type_subscript.add_child(Type.Any)
return type_subscript
def _parse_shape_structure(
self,
shape: StructureShape,
output: bool = False,
output_child: bool = False,
is_streaming: bool = False,
) -> FakeAnnotation:
if not shape.members.items():
return Type.DictStrAny if output_child else Type.MappingStrAny
required = shape.required_members
typed_dict_name = self._get_typed_dict_name(shape)
shape_type_stub = get_shape_type_stub(self.service_name,
typed_dict_name)
if shape_type_stub:
return shape_type_stub
typed_dict = TypeTypedDict(typed_dict_name)
if typed_dict.name in self._typed_dict_map:
old_typed_dict = self._typed_dict_map[typed_dict.name]
child_names = {i.name for i in old_typed_dict.children}
if output and "ResponseMetadata" in child_names:
return self._typed_dict_map[typed_dict.name]
if not output and "ResponseMetadata" not in child_names:
return self._typed_dict_map[typed_dict.name]
if output:
typed_dict.name = self._get_typed_dict_name(
shape, postfix="ResponseMetadata")
self.logger.debug(
f"Marking {typed_dict.name} as ResponseMetadataTypeDef")
else:
old_typed_dict.name = self._get_typed_dict_name(
shape, postfix="ResponseMetadata")
self._typed_dict_map[old_typed_dict.name] = old_typed_dict
self.logger.debug(
f"Marking {old_typed_dict.name} as ResponseMetadataTypeDef"
)
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,
output_child=output or output_child,
is_streaming=is_streaming,
),
attr_name in required,
)
if output:
self._make_output_typed_dict(typed_dict)
return typed_dict
def _make_output_typed_dict(self, typed_dict: TypeTypedDict) -> None:
for attribute in typed_dict.children:
attribute.required = True
child_names = {i.name for i in typed_dict.children}
if "ResponseMetadata" not in child_names:
typed_dict.add_attribute(
"ResponseMetadata",
self.response_metadata_typed_dict,
True,
)
def _parse_shape_list(self,
shape: ListShape,
output_child: bool = False) -> FakeAnnotation:
type_subscript = TypeSubscript(
Type.List) if output_child else TypeSubscript(Type.Sequence)
if shape.member:
type_subscript.add_child(
self.parse_shape(shape.member, output_child=output_child))
else:
type_subscript.add_child(Type.Any)
return type_subscript
def parse_shape(
self,
shape: Shape,
output: bool = False,
output_child: bool = False,
is_streaming: bool = False,
) -> FakeAnnotation:
"""
Parse any botocore shape to TypeAnnotation.
Arguments:
shape -- Botocore shape.
output -- Whether shape should use strict output types.
output_child -- Whether shape parent is marked as output.
is_streaming -- Whether shape should be streaming.
Returns:
TypeAnnotation or similar class.
"""
if not is_streaming:
is_streaming = "streaming" in shape.serialization and shape.serialization[
"streaming"]
if output or output_child:
is_streaming = self.proxy_operation_model._get_streaming_body(
shape) is not None # type: ignore
type_name = shape.type_name
if is_streaming and type_name == "blob":
type_name = "blob_streaming"
if output or output_child:
if type_name in self.OUTPUT_SHAPE_TYPE_MAP:
return self.OUTPUT_SHAPE_TYPE_MAP[type_name]
if type_name in self.SHAPE_TYPE_MAP:
return self.SHAPE_TYPE_MAP[type_name]
if isinstance(shape, StringShape):
return self._parse_shape_string(shape)
if isinstance(shape, MapShape):
return self._parse_shape_map(
shape,
output_child=output or output_child,
is_streaming=is_streaming,
)
if isinstance(shape, StructureShape):
return self._parse_shape_structure(
shape,
output=output,
output_child=output or output_child,
is_streaming=is_streaming,
)
if isinstance(shape, ListShape):
return self._parse_shape_list(shape,
output_child=output or output_child)
if self._resources_shape and shape.type_name in self._resources_shape[
"resources"]:
return AliasInternalImport(shape.type_name)
self.logger.warning(f"Unknown shape: {shape} {type_name}")
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:
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_wait_method(self, waiter_name: str) -> Method:
"""
Get Waiter `wait` method.
Arguments:
waiter_name -- Waiter name.
Returns:
Method.
"""
if not self._waiters_shape:
raise ShapeParserError("Waiter not found")
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:
shape_arguments = self._parse_arguments(
"Waiter", "wait", operation_name, operation_shape.input_shape)
shape_arguments.append(
Argument("WaiterConfig", waiter_config_type, Type.Ellipsis))
arguments.extend(self._get_kw_flags("wait", shape_arguments))
arguments.extend(shape_arguments)
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.Sequence, [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.Sequence, [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
@staticmethod
def _get_arg_from_target(target: str) -> str:
if "[" not in target:
return target
return target.split("[")[0]
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])
operation_shape = None
if "request" in action_shape:
operation_name = action_shape["request"]["operation"]
operation_shape = self._get_operation(operation_name)
skip_argument_names = {
self._get_arg_from_target(i["target"])
for i in action_shape["request"].get("params", {})
if i["source"] == "identifier"
}
if operation_shape.input_shape is not None:
shape_arguments = self._parse_arguments(
resource_name,
method_name,
operation_name,
operation_shape.input_shape,
exclude_names=skip_argument_names,
)
arguments.extend(
self._get_kw_flags(method_name, shape_arguments))
arguments.extend(shape_arguments)
if operation_shape.output_shape is not None and return_type is Type.none:
operation_return_type = self.parse_shape(
operation_shape.output_shape, output=True)
return_type = operation_return_type
method = Method(name=method_name,
arguments=arguments,
return_type=return_type)
if operation_shape and operation_shape.input_shape is not None:
method.request_type_annotation = method.get_request_type_annotation(
self._get_typed_dict_name(operation_shape.input_shape,
postfix=resource_name))
return method
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.
"""
result = Method(
name="filter",
arguments=[Argument("self", None)],
return_type=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:
shape_arguments = self._parse_arguments(
name,
result.name,
operation_name,
operation_model.input_shape,
optional_only=True,
)
result.arguments.extend(
self._get_kw_flags(result.name, shape_arguments))
result.arguments.extend(shape_arguments)
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(
name=batch_action.name,
arguments=[Argument("self", None)],
return_type=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:
shape_arguments = self._parse_arguments(
name,
batch_action.name,
operation_name,
operation_model.input_shape,
optional_only=True,
)
method.arguments.extend(
self._get_kw_flags(batch_action.name, shape_arguments))
method.arguments.extend(shape_arguments)
if operation_model.output_shape is not None:
item_return_type = self.parse_shape(
operation_model.output_shape, output=True)
return_type = TypeSubscript(Type.List, [item_return_type])
method.return_type = return_type
return result
def _parse_shape_map(
self,
shape: MapShape,
output_child: bool = False,
is_streaming: bool = False,
) -> FakeAnnotation:
type_subscript = TypeSubscript(
Type.Dict) if output_child else TypeSubscript(Type.Mapping)
if shape.key:
type_subscript.add_child(
self.parse_shape(shape.key,
output_child=output_child,
is_streaming=is_streaming))
else:
type_subscript.add_child(Type.str)
if shape.value:
type_subscript.add_child(
self.parse_shape(shape.value,
output_child=output_child,
is_streaming=is_streaming))
else:
type_subscript.add_child(Type.Any)
return type_subscript
from typing import Dict
from mypy_boto3_builder.import_helpers.import_string import ImportString
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]),
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.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,
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
