AWS SDK示例项目中的Cognito Python类型提示实践

在AWS SDK示例项目中，Python开发者经常需要与Amazon Cognito服务进行交互。然而，现有的代码示例中缺乏详细的类型提示，这给开发者带来了不小的困惑。本文将深入探讨如何为Cognito相关操作添加完善的类型提示，提升代码的可读性和可维护性。

Cognito客户端基础类型

首先，我们需要定义一个基础的CognitoClient类，它封装了与Cognito服务交互的基本配置：

from dataclasses import dataclass
from typing import Optional
import boto3

@dataclass(frozen=True, slots=True, kw_only=True)
class CognitoClient:
    """封装Amazon Cognito操作"""
    
    cognito_idp_client: boto3.client
    user_pool_id: str
    client_id: str
    client_secret: Optional[str] = None

这个类使用了Python的dataclass装饰器，并启用了slots优化内存使用。cognito_idp_client明确标注为boto3.client类型，其他参数也都添加了明确的类型注解。

用户属性与MFA配置

Cognito服务中的用户属性和MFA配置需要特定的数据结构。我们可以定义对应的数据类型：

@dataclass(frozen=True, slots=True, kw_only=True)
class Attribute:
    """表示Cognito用户属性"""
    
    Name: str
    Value: Optional[str]

@dataclass(frozen=True, slots=True, kw_only=True)
class MFAOption:
    """表示MFA配置选项"""
    
    AttributeName: Optional[str]
    DeliveryMedium: Optional[str]

用户数据类型

对于Cognito返回的用户数据，我们可以定义一个完整的User类：

from datetime import datetime

@dataclass(frozen=True, slots=True, kw_only=True)
class User:
    """表示Cognito用户对象"""
    
    Attributes: Optional[list[Attribute]]
    Enabled: Optional[bool]
    MFAOptions: Optional[list[MFAOption]]
    UserCreateDate: Optional[datetime]
    UserLastModifiedDate: Optional[datetime]
    Username: Optional[str]
    UserStatus: Optional[str]
    
    def __post_init__(self):
        if self.Attributes is not None:
            self.Attributes = [Attribute(**attr) for attr in self.Attributes]
        if self.MFAOptions is not None:
            self.MFAOptions = [MFAOption(**mfa) for mfa in self.MFAOptions]
        if self.UserCreateDate is not None:
            self.UserCreateDate = datetime.fromtimestamp(self.UserCreateDate)
        if self.UserLastModifiedDate is not None:
            self.UserLastModifiedDate = datetime.fromtimestamp(self.UserLastModifiedDate)

这个类不仅定义了所有可能的用户属性，还通过__post_init__方法实现了数据的自动转换，确保从API返回的数据能正确转换为Python对象。

核心操作的类型提示

用户注册

def sign_up_user(
    *,
    cognito_client: CognitoClient,
    user_name: str,
    password: str,
    user_email: str
) -> bool:
    """注册新用户"""
    try:
        kwargs = {
            "ClientId": cognito_client.client_id,
            "Username": user_name,
            "Password": password,
            "UserAttributes": [{"Name": "email", "Value": user_email}],
        }
        if cognito_client.client_secret is not None:
            kwargs["SecretHash"] = cognito_client.hash_from_user_name(user_name)
        response = cognito_client.cognito_idp_client.sign_up(**kwargs)
        return response["UserConfirmed"]
    except ClientError as err:
        # 错误处理逻辑

用户登录

def start_sign_in(
    *,
    cognito_client: CognitoClient,
    user_name: str,
    password: str
) -> dict[str, Any]:
    """初始化登录流程"""
    try:
        kwargs = {
            "UserPoolId": cognito_client.user_pool_id,
            "ClientId": cognito_client.client_id,
            "AuthFlow": "ADMIN_USER_PASSWORD_AUTH",
            "AuthParameters": {"USERNAME": user_name, "PASSWORD": password},
        }
        if cognito_client.client_secret is not None:
            kwargs["AuthParameters"]["SECRET_HASH"] = cognito_client.hash_from_user_name(user_name)
        response = cognito_client.cognito_idp_client.admin_initiate_auth(**kwargs)
        response.pop("ResponseMetadata", None)
        return response

MFA验证

def respond_to_mfa_challenge(
    *,
    cognito_client: CognitoClient,
    user_name: str,
    session: str,
    mfa_code: