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 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_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 __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_method(parent_name: str, name: str, method: FunctionType, 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) 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) return Method(name=name, arguments=arguments, return_type=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_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(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_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_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_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 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(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_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_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_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 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 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 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 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 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
def parse_collections( parent_name: str, resource: Boto3ServiceResource, service_name: ServiceName, shape_parser: ShapeParser, ) -> List[Collection]: """ Extract collections from boto3 resource. Arguments: resource -- boto3 service resource. Returns: A list of Collection structures. """ result: List[Collection] = [] for collection in resource.meta.resource_model.collections: if not collection.resource: continue object_class_name = collection.resource.type collection_record = Collection( name=f"{parent_name}{get_class_prefix(collection.name)}Collection", parent_name=parent_name, attribute_name=collection.name, service_name=service_name, type_annotation=InternalImport(collection.name), object_class_name=object_class_name, ) self_type = InternalImport(collection_record.name, stringify=True) collection_record.methods.append( Method( name="all", arguments=[Argument("self", None)], return_type=self_type, docstring=("Get all items from the collection, optionally" " with a custom page size and item count limit."), )) filter_method = shape_parser.get_collection_filter_method( collection_record.name, collection, self_type) filter_method.docstring = ( "Get items from the collection, passing keyword arguments along" " as parameters to the underlying service operation, which are" " typically used to filter the results.") filter_method.type_ignore = True collection_record.methods.append(filter_method) batch_methods = shape_parser.get_collection_batch_methods( collection_record.name, collection) for batch_method in batch_methods: batch_method.docstring = "Batch method." collection_record.methods.append(batch_method) collection_record.methods.append( Method( "limit", [Argument("self", None), Argument("count", Type.int)], self_type, docstring=f"Return at most this many {object_class_name}s.", )) collection_record.methods.append( Method( "page_size", [Argument("self", None), Argument("count", Type.int)], self_type, docstring= f"Fetch at most this many {object_class_name}s per service request.", )) collection_record.methods.append( Method( "pages", [Argument("self", None)], TypeSubscript( Type.Iterator, [ TypeSubscript(Type.List, [InternalImport(name=object_class_name)]) ], ), docstring= f"A generator which yields pages of {object_class_name}s.", )) collection_record.methods.append( Method( "__iter__", [Argument("self", None)], TypeSubscript(Type.Iterator, [InternalImport(name=object_class_name)]), docstring=f"A generator which yields {object_class_name}s.", )) result.append(collection_record) return result
def parse_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
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
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