如何为泛型容器的search方法实现类型提示及类SQLAlchemy搜索语法

一、为column参数添加类型提示（支持自动补全） #

可以结合Pydantic模型的model_fields和Literal类型，实现对column参数的类型约束，让IDE自动补全合法字段名：

from typing import Any, Generic, TypeVar, Literal
from pydantic import BaseModel

T = TypeVar("T", bound=BaseModel)

class Container(Generic[T]):
    def __init__(self, values: list[T]) -> None:
        self._values = values

    def search(self, column: Literal[tuple(T.model_fields.keys())], value: Any) -> list[T]:
        return [r for r in self._values if getattr(r, column) == value]

class SomeData(BaseModel):
    col1: int
    col2: int

当你在IDE中调用c.search(column=...)时，会自动弹出col1、col2的补全选项，传入非法列名会触发类型检查错误。

二、实现类似SQLAlchemy的查询语法 #

要实现c.search(col1 == 1)或c.search(c.col1 == 1)的语法，需要自定义字段表达式类，让字段支持比较操作并生成查询条件，再让search方法解析这些条件。

1. 自定义字段表达式类 #

首先定义一个FieldExpr类，用来封装字段名和比较操作：

class FieldExpr:
    def __init__(self, field_name: str):
        self.field_name = field_name

    def __eq__(self, other) -> tuple[str, str, Any]:
        return ("eq", self.field_name, other)
    
    # 可扩展其他操作符，比如大于、小于
    def __gt__(self, other) -> tuple[str, str, Any]:
        return ("gt", self.field_name, other)
    
    def __lt__(self, other) -> tuple[str, str, Any]:
        return ("lt", self.field_name, other)

2. 让模型或容器支持字段表达式 #

方式一：通过模型类调用字段（SomeData.col1 == 1）

利用元类自动为模型的每个字段生成对应的类方法，返回FieldExpr实例：

from typing import Any, Generic, TypeVar
from pydantic import BaseModel

class ModelMeta(type(BaseModel)):
    def __new__(cls, name, bases, attrs):
        new_cls = super().__new__(cls, name, bases, attrs)
        # 为每个模型字段自动生成类方法
        for field_name in new_cls.model_fields:
            def make_field_expr(field=field_name):
                @classmethod
                def field_expr(cls):
                    return FieldExpr(field)
                return field_expr
            setattr(new_cls, field_name, make_field_expr())
        return new_cls

# 模型继承该元类
class SomeData(BaseModel, metaclass=ModelMeta):
    col1: int
    col2: int

方式二：通过容器实例调用字段（c.col1 == 1）

在容器初始化时，动态为实例添加字段属性，返回FieldExpr实例：

class Container(Generic[T]):
    def __init__(self, values: list[T]) -> None:
        self._values = values
        # 动态添加对应模型字段的属性
        for field_name in T.model_fields:
            setattr(self, field_name, FieldExpr(field_name))

3. 改造search方法解析条件 #

更新Container的search方法，支持解析FieldExpr生成的条件元组：

class Container(Generic[T]):
    def __init__(self, values: list[T]) -> None:
        self._values = values
        # 方式二需要的动态属性
        for field_name in T.model_fields:
            setattr(self, field_name, FieldExpr(field_name))

    def search(self, condition: tuple[str, str, Any]) -> list[T]:
        op, field, value = condition
        if op == "eq":
            return [r for r in self._values if getattr(r, field) == value]
        elif op == "gt":
            return [r for r in self._values if getattr(r, field) > value]
        elif op == "lt":
            return [r for r in self._values if getattr(r, field) < value]
        # 可扩展更多操作逻辑
        return []

使用示例 #

# 方式一：通过模型类调用字段
c = Container(values=[SomeData(col1=1, col2=2), SomeData(col1=10, col2=20)])
print(c.search(SomeData.col1() == 1)) # 返回 [SomeData(col1=1, col2=2)]

# 方式二：通过容器实例调用字段
print(c.search(c.col2 > 5)) # 返回 [SomeData(col1=10, col2=20)]

内容的提问来源于stack exchange，提问作者Guillaume