将SQL查询初始化为dataclass时如何避免get_fields函数硬编码问题

我们的目标是基于对字段的人工分类得到SQL中需要查询的列，示例代码如下：

from dataclasses import dataclass

@dataclass
class Q:
    fruits = ('taste', 'color', 'shine')
    vege = ('is_green', 'color', 'juiceable')
    meat = ('is_white', 'is_dark', 'is_red', 'meat_name')
    carbs = ('is_refined', 'gi')

    @staticmethod
    def get_fields(self, add_fruits=False, add_vege=False, add_meat=False, add_carbs=False):
         fields = set()
         fields.update(self.fruits) if add_fruits else None
         fields.update(self.vege) if add_vege else None
         fields.update(self.meat) if add_meat else None
         fields.update(self.carbs) if add_carbs else None
         return ",".join(sorted(fields))

print(f"SELECT {Q.get_fields(add_fruits=True, add_vege=True)} from food_features")

注：原示例代码存在两处可优化的小问题：一是@staticmethod不需要接收self参数，要访问类属性建议改用@classmethod；二是原代码对join后的字符串做sorted会把字符串拆成单个字符列表，此处已修正为仅对字段列表排序后再拼接。

当Q类中存在大量字段分类时，如何避免在.get_fields函数中硬编码如下逻辑？

fields.update(self._field_name_) if add_field_name else None 

解决方案 #

方案1：自动匹配add_前缀参数（兼容原有调用方式） #

无需修改现有调用逻辑，新增分类时仅需要在Q类中新增对应属性即可，不需要改动get_fields代码：

from dataclasses import dataclass

@dataclass
class Q:
    fruits = ('taste', 'color', 'shine')
    vege = ('is_green', 'color', 'juiceable')
    meat = ('is_white', 'is_dark', 'is_red', 'meat_name')
    carbs = ('is_refined', 'gi')
    # 后续新增分类直接在这里加属性即可，不用改get_fields
    # snack = ('is_salty', 'is_sweet', 'calorie')

    @classmethod
    def get_fields(cls, **kwargs):
        fields = set()
        # 遍历所有关键字参数，匹配add_前缀的参数
        for param_name, is_add in kwargs.items():
            if param_name.startswith('add_') and is_add:
                # 去掉add_前缀得到对应的分类属性名
                category = param_name[4:]
                # 从类中取出对应分类的字段列表加入集合
                if hasattr(cls, category):
                    fields.update(getattr(cls, category))
        return ",".join(sorted(fields))

# 原有调用方式完全兼容
print(f"SELECT {Q.get_fields(add_fruits=True, add_vege=True)} from food_features")
# 新增分类后直接传对应参数即可
# print(f"SELECT {Q.get_fields(add_snack=True, add_meat=True)} from food_features")

方案2：直接传入分类列表（调用更简洁） #

如果可以调整调用方式，直接传入需要的分类名称列表，代码逻辑更清晰易懂：

from dataclasses import dataclass

@dataclass
class Q:
    fruits = ('taste', 'color', 'shine')
    vege = ('is_green', 'color', 'juiceable')
    meat = ('is_white', 'is_dark', 'is_red', 'meat_name')
    carbs = ('is_refined', 'gi')

    @classmethod
    def get_fields(cls, categories: list):
        fields = set()
        for cate in categories:
            if hasattr(cls, cate):
                fields.update(getattr(cls, cate))
        return ",".join(sorted(fields))

# 调用时直接传入需要的分类列表
print(f"SELECT {Q.get_fields(categories=['fruits', 'vege'])} from food_features")

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