diff --git a/web/collectors/__init__.py b/web/collectors/__init__.py index 37660cd..fb6b5d7 100644 --- a/web/collectors/__init__.py +++ b/web/collectors/__init__.py @@ -16,7 +16,7 @@ class Collector: self.force = True # todo-2 已发布的价格不在覆盖计算 self.digit_map = {} # 获取所有材料信息 - self.materials = Material.query.all() + self.materials = Material.list(type=1) self.material_codes = [m.code for m in self.materials if m.code] # 缓存材料税率信息 self.tax_map = {m.code: m.tax for m in self.materials if m.code} diff --git a/web/commons/models/material.py b/web/commons/models/material.py index 34da408..ccadcae 100644 --- a/web/commons/models/material.py +++ b/web/commons/models/material.py @@ -35,6 +35,13 @@ class Material(db.Model, Model, OperationTrackMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + @classmethod + def list(cls, type=None): + query = cls.query + if type: + query = query.filter(cls.type == type) + return query.all() + @classmethod def get_by_key(cls, code): return cls.query.filter(cls.code == code).one_or_none() diff --git a/web/tasks/once/calculate.py b/web/tasks/once/calculate.py index 92b5f21..17199f3 100644 --- a/web/tasks/once/calculate.py +++ b/web/tasks/once/calculate.py @@ -12,7 +12,7 @@ def calculate(year=2023, month=8): # 清空当月数据 PriceResult.clean(year, month) # 从材料表获取趋势表所需的材料信息 - for material in Material.query.all(): + for material in Material.list(type=1): calculator = Calculator(year=year, month=month) calculator.name = material.name calculator.material_id = material.id