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