class TestFunction: def setup_method(self): self.function = Function( name="name", arguments=[ Argument(self, None), Argument("my_str", Type.str, TypeConstant("test")), Argument("lst", Type.ListAny), ], decorators=[Type.Any], return_type=Type.none, body_lines=["line1", "line2"], ) def test_init(self) -> None: assert self.function.name == "name" assert len(self.function.arguments) == 3 assert self.function.body == "line1\nline2" def test_get_types(self) -> None: assert self.function.get_types() == { Type.Any, Type.none, Type.List, Type.str, TypeConstant("test"), } def test_get_required_import_records(self) -> None: assert self.function.get_required_import_records() == { ImportRecord(ImportString("typing"), "Any"), ImportRecord(ImportString("typing"), "List"), }
def setup_method(self): self.function = Function( name="name", arguments=[ Argument(self, None), Argument("my_str", Type.str, TypeConstant("test")), Argument("lst", Type.ListAny), ], decorators=[Type.Any], return_type=Type.none, body_lines=["line1", "line2"], )
def 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