Python 高级类型提示实战：Protocol、Generic 与 TypeVar 的深度用法

1
from abc import ABC, abstractmethod
2

3
class Drawable(ABC):
4
    @abstractmethod
5
    def draw(self) -> None: ...
6

7
class Circle(Drawable):
8
    def draw(self) -> None:
9
        print("画圆")
10

11
class Square(Drawable):
12
    def draw(self) -> None:
13
        print("画方块")
14

15
# ❌ 这个类没有继承 Drawable，但它有 draw() 方法
16
class Triangle:
17
    def draw(self) -> None:
18
        print("画三角形")
19

20
def render(shape: Drawable) -> None:
21
    shape.draw()
22

23
render(Circle())   # ✅ 可以
24
render(Square())   # ✅ 可以
25
render(Triangle()) # ❌ mypy 报错：Triangle 不是 Drawable 的子类

1
from typing import Protocol, runtime_checkable
2

3
class HasDraw(Protocol):
4
    def draw(self) -> None: ...
5

6
class Circle:
7
    def draw(self) -> None:
8
        print("画圆")
9

10
class Square:
11
    def draw(self) -> None:
12
        print("画方块")
13

14
# Triangle 不需要继承任何东西
15
class Triangle:
16
    def draw(self) -> None:
17
        print("画三角形")
18

19
def render(shape: HasDraw) -> None:
20
    shape.draw()
21

22
render(Circle())    # ✅
23
render(Square())    # ✅
24
render(Triangle())  # ✅ mypy 也能通过了！

1
from typing import Protocol, runtime_checkable
2

3
@runtime_checkable
4
class Serializable(Protocol):
5
    def to_json(self) -> str: ...
6
    @classmethod
7
    def from_json(cls, data: str) -> "Serializable": ...
8

9
class User:
10
    def __init__(self, name: str, age: int) -> None:
11
        self.name = name
12
        self.age = age
13

14
    def to_json(self) -> str:
15
        import json
16
        return json.dumps({"name": self.name, "age": self.age})
17

18
    @classmethod
19
    def from_json(cls, data: str) -> "User":
20
        import json
21
        d = json.loads(data)
22
        return cls(d["name"], d["age"])
23

24
user = User("张三", 28)
25
print(isinstance(user, Serializable))  # True！

1
from typing import Protocol, TypeVar, Generic
2

3
T = TypeVar("T")
4

5
class Repository(Protocol[T]):
6
    def get(self, id: int) -> T | None: ...
7
    def save(self, entity: T) -> None: ...
8
    def delete(self, id: int) -> bool: ...
9

10
# 实现：用户仓库
11
class User:
12
    def __init__(self, id: int, name: str) -> None:
13
        self.id = id
14
        self.name = str
15

16
class UserRepository:
17
    def __init__(self) -> None:
18
        self._store: dict[int, User] = {}
19

20
    def get(self, id: int) -> User | None:
21
        return self._store.get(id)
22

23
    def save(self, entity: User) -> None:
24
        self._store[entity.id] = entity
25

26
    def delete(self, id: int) -> bool:
27
        return self._store.pop(id, None) is not None
28

29
def process_repo(repo: Repository[User]) -> None:
30
    user = repo.get(1)
31
    if user:
32
        repo.save(user)
33

34
repo = UserRepository()
35
process_repo(repo)  # ✅ mypy 自动验证类型匹配

1
from typing import Protocol, TypeVar, Generic
2

3
T_in = TypeVar("T_in")
4
T_out = TypeVar("T_out")
5

6
class Transformer(Protocol[T_in, T_out]):
7
    """数据转换器协议：接收 T_in，产出 T_out"""
8
    def transform(self, data: T_in) -> T_out: ...
9

10
class Pipeline(Generic[T_in, T_out]):
11
    def __init__(self) -> None:
12
        self._steps: list[Transformer] = []
13

14
    def add(self, step: Transformer) -> None:
15
        self._steps.append(step)
16

17
    def execute(self, data: T_in) -> T_out:
18
        result: object = data
19
        for step in self._steps:
20
            result = step.transform(result)
21
        return result  # type: ignore[return-value]
22

23
# 具体转换器
24
class StringToUpper:
25
    def transform(self, data: str) -> str:
26
        return data.upper()
27

28
class StringLength:
29
    def transform(self, data: str) -> int:
30
        return len(data)
31

32
pipeline: Pipeline[str, int] = Pipeline()
33
pipeline.add(StringToUpper())
34
pipeline.add(StringLength())
35
result = pipeline.execute("hello world")
36
print(result)  # 11

1
from typing import TypeVar, Generic
2

3
class Animal: ...
4
class Dog(Animal): ...
5

6
T_co = TypeVar("T_co", covariant=True)   # 协变
7
T_contra = TypeVar("T_contra", contravariant=True)  # 逆变
8

9
class Producer(Generic[T_co]):
10
    """生产者：产出 T_co"""
11
    def produce(self) -> T_co: ...
12

13
class Consumer(Generic[T_contra]):
14
    """消费者：接收 T_contra"""
15
    def consume(self, item: T_contra) -> None: ...
16

17
# 协变：Producer[Dog] 是 Producer[Animal] 的子类型
18
dog_producer: Producer[Dog] = Producer()
19
animal_producer: Producer[Animal] = dog_producer  # ✅ 协变安全
20

21
# 逆变：Consumer[Animal] 是 Consumer[Dog] 的子类型
22
animal_consumer: Consumer[Animal] = Consumer()
23
dog_consumer: Consumer[Dog] = animal_consumer  # ✅ 逆变安全

1
from typing import Protocol, TypeVar, Generic
2

3
T_co = TypeVar("T_co", covariant=True)
4

5
class Event(Protocol[T_co]):
6
    """事件协议：携带 T_co 类型的数据"""
7
    @property
8
    def data(self) -> T_co: ...
9
    @property
10
    def timestamp(self) -> float: ...
11

12
T_contra = TypeVar("T_contra", contravariant=True)
13

14
class EventHandler(Protocol[T_contra]):
15
    """事件处理器：能处理 T_contra 类型的事件"""
16
    def handle(self, event: T_contra) -> None: ...
17

18
class EventBus:
19
    def __init__(self) -> None:
20
        self._handlers: dict[type, list[EventHandler]] = {}
21

22
    def register(self, event_type: type, handler: EventHandler) -> None:
23
        self._handlers.setdefault(event_type, []).append(handler)
24

25
    def dispatch(self, event: Event) -> None:
26
        handlers = self._handlers.get(type(event), [])
27
        for handler in handlers:
28
            handler.handle(event)

1
from typing import Generic, TypeVar, Callable
2
import time
3

4
K = TypeVar("K")
5
V = TypeVar("V")
6

7
class LRUCache(Generic[K, V]):
8
    def __init__(self, capacity: int, ttl: float = 300.0) -> None:
9
        self._capacity = capacity
10
        self._ttl = ttl
11
        self._store: dict[K, tuple[V, float]] = {}
12
        self._order: list[K] = []
13

14
    def get(self, key: K) -> V | None:
15
        if key not in self._store:
16
            return None
17

18
        value, timestamp = self._store[key]
19
        if time.time() - timestamp > self._ttl:
20
            del self._store[key]
21
            self._order.remove(key)
22
            return None
23

24
        # 移到末尾（最近使用）
25
        self._order.remove(key)
26
        self._order.append(key)
27
        return value
28

29
    def put(self, key: K, value: V) -> None:
30
        if key in self._store:
31
            self._order.remove(key)
32
        elif len(self._store) >= self._capacity:
33
            # 淘汰最旧的
34
            oldest = self._order.pop(0)
35
            del self._store[oldest]
36

37
        self._store[key] = (value, time.time())
38
        self._order.append(key)
39

40
    def __len__(self) -> int:
41
        return len(self._store)
42

43
# 使用：类型完全自动推导
44
cache: LRUCache[str, int] = LRUCache(capacity=3, ttl=60.0)
45
cache.put("a", 1)
46
cache.put("b", 2)
47
result = cache.get("a")  # result 类型推导为 int | None

1
from typing import Generic, TypeVar, Callable
2

3
T = TypeVar("T")
4
E = TypeVar("E")
5

6
class Ok(Generic[T]):
7
    def __init__(self, value: T) -> None:
8
        self._value = value
9

10
    @property
11
    def value(self) -> T:
12
        return self._value
13

14
    @property
15
    def is_ok(self) -> bool:
16
        return True
17

18
    @property
19
    def is_err(self) -> bool:
20
        return False
21

22
    def map(self, fn: Callable[[T], T]) -> "Ok[T]":
23
        return Ok(fn(self._value))
24

25
    def map_err(self, fn: Callable[[E], E]) -> "Ok[T]":
26
        return self  # Ok 状态下不转换错误
27

28
    def unwrap(self) -> T:
29
        return self._value
30

31
class Err(Generic[E]):
32
    def __init__(self, error: E) -> None:
33
        self._error = error
34

35
    @property
36
    def error(self) -> E:
37
        return self._error
38

39
    @property
40
    def is_ok(self) -> bool:
41
        return False
42

43
    @property
44
    def is_err(self) -> bool:
45
        return True
46

47
    def map(self, fn: Callable) -> "Err[E]":
48
        return self  # Err 状态下不转换值
49

50
    def map_err(self, fn: Callable[[E], E]) -> "Err[E]":
51
        return Err(fn(self._error))
52

53
    def unwrap(self) -> None:
54
        raise RuntimeError(f"Called unwrap() on Err: {self._error}")
55

56
Result = Ok[T] | Err[E]
57

58
# 使用示例
59
def divide(a: float, b: float) -> Result[float, str]:
60
    if b == 0:
61
        return Err("除数不能为零")
62
    return Ok(a / b)
63

64
result = divide(10, 3)
65
if result.is_ok:
66
    print(f"结果: {result.unwrap()}")  # 结果: 3.3333333333333335
67
else:
68
    print(f"错误: {result.error}")
69

70
result2 = divide(10, 0)
71
if result2.is_ok:
72
    print(f"结果: {result2.unwrap()}")
73
else:
74
    print(f"错误: {result2.error}")  # 错误: 除数不能为零

1
# Python 3.12+（不需要 from typing import TypeVar, Generic）
2

3
class Stack[T]:
4
    def __init__(self) -> None:
5
        self._items: list[T] = []
6

7
    def push(self, item: T) -> None:
8
        self._items.append(item)
9

10
    def pop(self) -> T:
11
        if not self._items:
12
            raise IndexError("Stack is empty")
13
        return self._items.pop()
14

15
class Result[T, E]:
16
    def __init__(self, ok: bool, value: T | None = None, error: E | None = None) -> None:
17
        self.ok = ok
18
        self.value = value
19
        self.error = error
20

21
# 泛型函数也可以直接写
22
def first[T](items: list[T]) -> T:
23
    if not items:
24
        raise IndexError("List is empty")
25
    return items[0]

1
from typing import Callable, TypeVar
2

3
T = TypeVar("T")
4

5
# ❌ 传统方式：装饰器会丢失被装饰函数的参数类型
6
def logged(func: Callable[..., T]) -> Callable[..., T]:
7
    def wrapper(*args, **kwargs):
8
        print(f"Calling {func.__name__}")
9
        return func(*args, **kwargs)
10
    return wrapper
11

12
@logged
13
def add(a: int, b: int) -> int:
14
    return a + b
15

16
# mypy 无法推导 add 的参数类型
17
# add 被推导为 Callable[..., int]，丢失了 (int, int) 的信息

1
from typing import Callable, TypeVar, ParamSpec
2

3
P = ParamSpec("P")
4
T = TypeVar("T")
5

6
def logged(func: Callable[P, T]) -> Callable[P, T]:
7
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
8
        print(f"Calling {func.__name__}")
9
        return func(*args, **kwargs)
10
    return wrapper
11

12
@logged
13
def add(a: int, b: int) -> int:
14
    return a + b
15

16
# ✅ mypy 正确推导：add 仍然是 (int, int) -> int
17
result = add(1, 2)  # 类型推导为 int
18
# add("hello", 2)   # ❌ mypy 会报错！

1
from typing import Callable, ParamSpec, TypeVar
2
import time
3

4
P = ParamSpec("P")
5
T = TypeVar("T")
6

7
def retry(max_attempts: int = 3, delay: float = 1.0) -> Callable[[Callable[P, T]], Callable[P, T]]:
8
    def decorator(func: Callable[P, T]) -> Callable[P, T]:
9
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
10
            last_exception: Exception | None = None
11
            for attempt in range(max_attempts):
12
                try:
13
                    return func(*args, **kwargs)
14
                except Exception as e:
15
                    last_exception = e
16
                    print(f"Attempt {attempt + 1}/{max_attempts} failed: {e}")
17
                    if attempt < max_attempts - 1:
18
                        time.sleep(delay)
19
            raise last_exception  # type: ignore[misc]
20
        return wrapper
21
    return decorator
22

23
@retry(max_attempts=3, delay=0.5)
24
def unstable_api(x: int) -> str:
25
    import random
26
    if random.random() < 0.7:
27
        raise ConnectionError("网络不稳定")
28
    return f"success: {x}"
29

30
result = unstable_api(42)  # 类型推导为 str

1
from typing import Generic, TypeVar, TypeVarTuple
2

3
Ts = TypeVarTuple("Ts")
4
T = TypeVar("T")
5

6
class Tensor(Generic[*Ts]):
7
    """N 维张量：每个维度可以有不同类型"""
8
    def __init__(self, data: object, shape: tuple[*Ts]) -> None:
9
        self.data = data
10
        self.shape = shape
11

12
# 2D 张量
13
tensor_2d: Tensor[int, str] = Tensor([[1, 2], [3, 4]], shape=(2, 2))
14

15
# 3D 张量
16
tensor_3d: Tensor[int, str, float] = Tensor(
17
    [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
18
    shape=(2, 2, 2)
19
)

1
from typing import Generic, TypeVar, TypeVarTuple, Callable, cast
2

3
T = TypeVar("T")
4
Ts = TypeVarTuple("Ts")
5

6
class Pipe(Generic[T, *Ts]):
7
    """类型安全的函数管道"""
8
    def __init__(self, initial: T) -> None:
9
        self._value: object = initial
10

11
    def pipe(self, fn: Callable[[T], T]) -> "Pipe[T, *Ts]":
12
        self._value = fn(self._value)  # type: ignore[assignment]
13
        return self
14

15
    def result(self) -> T:
16
        return cast(T, self._value)
17

18
def add_one(x: int) -> int:
19
    return x + 1
20

21
def double(x: int) -> int:
22
    return x * 2
23

24
def to_str(x: int) -> str:
25
    return str(x)
26

27
# 使用管道
28
result = (
29
    Pipe(10)
30
    .pipe(add_one)    # 11
31
    .pipe(double)     # 22
32
    .pipe(to_str)     # "22"
33
    .result()
34
)
35
print(result)  # "22"

1
from typing import TypeVar
2

3
T = TypeVar("T", bound="Builder")
4

5
class Builder:
6
    def __init__(self) -> None:
7
        self._parts: list[str] = []
8

9
    def add_part(self, part: str) -> T:  # ❌ 子类的类型会丢失
10
        self._parts.append(part)
11
        return self  # type: ignore[return-value]
12

13
class AdvancedBuilder(Builder):
14
    def add_special(self) -> "AdvancedBuilder":
15
        self._parts.append("special")
16
        return self
17

18
# 问题：
19
# Builder().add_part("a").add_part("b")  # 返回类型是 Builder
20
# AdvancedBuilder().add_part("a")         # 返回类型是 Builder，不是 AdvancedBuilder！

1
from typing import Self
2

3
class Builder:
4
    def __init__(self) -> None:
5
        self._parts: list[str] = []
6

7
    def add_part(self, part: str) -> Self:
8
        self._parts.append(part)
9
        return self
10

11
    def build(self) -> str:
12
        return " + ".join(self._parts)
13

14
class AdvancedBuilder(Builder):
15
    def add_special(self) -> Self:
16
        self._parts.append("special")
17
        return self
18

19
# ✅ 现在类型正确推导
20
ab = AdvancedBuilder()
21
result = ab.add_part("a").add_special().add_part("b").build()
22
print(result)  # "a + special + b"