def add_contextmanager_methods(self) -> None:
"""
Add contextmanager methods.
"""
self.package.client.methods.append(
Method(
"__aenter__",
[
Argument("self", None),
],
return_type=InternalImport(self.package.client.name),
is_async=True,
docstring=self.package.client.docstring,
)
)
self.package.client.methods.append(
Method(
"__aexit__",
[
Argument("self", None),
Argument("exc_type", Type.Any),
Argument("exc_val", Type.Any),
Argument("exc_tb", Type.Any),
],
return_type=Type.Any,
is_async=True,
docstring=self.package.client.docstring,
)
)
def __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 _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 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_arguments_from_argspec(func: MethodType) -> list[Argument]:
arguments: list[Argument] = []
argspec = inspect.getfullargspec(func)
for argument_name in argspec.args:
if argument_name == "factory_self":
argument_name = "self"
type_annotation: TypeAnnotation | None = Type.Any
if not arguments and argument_name in ("self", "cls"):
type_annotation = None
arguments.append(Argument(argument_name, type_annotation))
if argspec.defaults:
for index, default_value in enumerate(argspec.defaults):
argument_index = len(arguments) - len(argspec.defaults) + index
default_typed_value = TypeConstant(default_value)
if default_typed_value.is_none():
default_typed_value = Type.Ellipsis
arguments[argument_index].default = default_typed_value
if argspec.varargs:
arguments.append(Argument(argspec.varargs, Type.Any, prefix="*"))
for argument_name in argspec.kwonlyargs:
arguments.append(Argument(argument_name, Type.Any))
if argspec.kwonlydefaults:
for argument_name, default_value in argspec.kwonlydefaults.items():
for argument in arguments:
if argument.name != argument_name:
continue
argument.default = TypeConstant(default_value)
break
if argspec.varkw:
arguments.append(Argument(argspec.varkw, Type.Any, prefix="**"))
return arguments
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_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_arguments_from_argspec(func: FunctionType) -> List[Argument]:
arguments: List[Argument] = []
argspec = inspect.getfullargspec(func)
for argument_name in argspec.args:
if argument_name == "factory_self":
argument_name = "self"
type_annotation: Optional[TypeAnnotation] = Type.Any
if not arguments and argument_name in ("self", "cls"):
type_annotation = None
arguments.append(Argument(argument_name, type_annotation))
if argspec.defaults:
for index, default_value in enumerate(argspec.defaults):
argument_index = len(arguments) - len(argspec.defaults) + index
arguments[argument_index].default = TypeConstant(default_value)
if argspec.varargs:
arguments.append(Argument(argspec.varargs, Type.Any, prefix="*"))
for argument_name in argspec.kwonlyargs:
arguments.append(Argument(argument_name, Type.Any))
if argspec.kwonlydefaults:
for argument_name, default_value in argspec.kwonlydefaults:
for argument in arguments:
if argument.name != argument_name:
continue
argument.default = TypeConstant(default_value)
break
if argspec.varkw:
arguments.append(Argument(argspec.varkw, Type.Any, prefix="**"))
return arguments
def setup_method(self):
self.function = Function(
name="name",
arguments=[
Argument(self, None),
Argument("my_str", Type.str, TypeConstant("test")),
Argument("lst", Type.ListAny),
],
decorators=[Type.Any],
return_type=Type.none,
body_lines=["line1", "line2"],
)
def get_client_method(self) -> Method:
return Method(
name="get_paginator",
decorators=[Type.overload],
docstring=self.docstring,
arguments=[
Argument("self", None),
Argument("operation_name", TypeLiteral(self.operation_name)),
],
return_type=InternalImport(
self.name, module_name=ServiceModuleName.paginator),
)
def get_client_method(self) -> Method:
return Method(
name="get_waiter",
decorators=[Type.overload],
docstring=self.docstring,
arguments=[
Argument("self", None),
Argument("waiter_name", TypeLiteral(self.waiter_name)),
],
return_type=ExternalImport(
source=ImportString(self.service_name.module_name,
ServiceModuleName.waiter.value),
name=self.name,
),
)
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_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"]))
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:
return_type = self._parse_shape(operation_shape.output_shape)
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_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
def setup_method(self):
self.class_record = ClassRecord(
name="Name",
methods=[
Method(
name="my_method",
arguments=[
Argument("self", None),
Argument("my_str", Type.str, TypeConstant("test")),
Argument("lst", Type.ListAny),
],
return_type=Type.none,
)
],
attributes=[Attribute("attr", Type.Any, Type.none)],
bases=[Type.Any],
use_alias=True,
)
def _find_argument_or_append(self, name: str) -> Argument:
if name in self.arguments_map:
return self.arguments_map[name]
for key in list(self.arguments_map):
if key.startswith("*"):
del self.arguments_map[key]
self.arguments_map[name] = Argument(name, Type.Any, Type.none)
return self.arguments_map[name]
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,
),
)
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 test_init(self) -> None:
method = Method(
"my_method",
[
Argument("self", None),
Argument("second", Type.str),
Argument("other", Type.DictStrAny),
],
Type.none,
)
assert len(method.call_arguments) == 2
method = Method(
"my_method",
[
Argument("second", Type.str),
Argument("other", Type.DictStrAny),
],
Type.none,
)
assert len(method.call_arguments) == 2
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
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.none
if optional_only and argument.required:
continue
result.append(argument)
result.sort(key=lambda x: not x.required)
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])
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_exceptions_property(self) -> Method:
"""
Generate Client exceptions property.
"""
return Method(
name="exceptions",
decorators=[TypeClass(property)],
arguments=[
Argument("self", None),
],
return_type=InternalImport(
name=self.exceptions_class.name,
module_name=ServiceModuleName.client,
service_name=self.service_name,
stringify=False,
),
docstring=f"{self.name} exceptions.",
)
def parse_method(
parent_name: str, name: str, method: MethodType, service_name: ServiceName
) -> Method:
"""
Parse method to a structure.
Arguments:
parent_name -- Parent class name.
method -- Inspect method.
Returns:
Method structure.
"""
logger = get_logger()
docstring = textwrap.dedent(inspect.getdoc(method) or "")
method_name = f"{parent_name}.{name}"
logger.debug(f"Slow parsing of {method_name}: {len(docstring)} chars")
prefix = f"{get_class_prefix(parent_name)}{get_class_prefix(name)}"
arg_spec_parser = ArgSpecParser(prefix, service_name)
arguments = get_method_arguments_stub(service_name, parent_name, name)
if arguments is None:
arguments = arg_spec_parser.get_arguments(parent_name, name, method)
docstring_parser = DocstringParser(service_name, parent_name, name, arguments)
arguments = docstring_parser.get_arguments(docstring)
# do not add kwonly flag to resource generators
if len(arguments) > 1 and not name[0].isupper():
arguments.insert(1, Argument.kwflag())
return_type = arg_spec_parser.get_return_type(parent_name, name)
if return_type is None:
return_type = DocstringParser(service_name, parent_name, name, []).get_return_type(
docstring
)
result = Method(name=name, arguments=arguments, return_type=return_type)
result.request_type_annotation = result.get_request_type_annotation(
f"{parent_name}{get_class_prefix(name)}RequestTypeDef"
)
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 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.
Returns:
List of Method records.
"""
result = []
for batch_action in collection.batch_actions:
method = Method(
batch_action.name,
[Argument("cls", None)],
Type.none,
decorators=[Type.classmethod],
)
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 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_collection_filter_method(self, name: str,
collection: Collection) -> Method:
"""
Get `filter` classmethod for Resource collection.
Arguments:
name -- Collection record name.
collection -- Boto3 Collection.
Returns:
Filter Method record.
"""
arguments: List[Argument] = [Argument("cls", None)]
result = Method(
"filter",
arguments,
InternalImport(name=name, service_name=self.service_name),
decorators=[Type.classmethod],
)
if not collection.request:
return result
operation_name = collection.request.operation
operation_model = self._get_operation(operation_name)
if operation_model.input_shape is not None:
for argument in self._parse_arguments(
name,
result.name,
operation_name,
operation_model.input_shape,
):
if argument.required:
continue
arguments.append(argument)
return result
def 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