diff --git a/backend/app/modelling/property_filters.py b/backend/app/modelling/property_filters.py index 414f16889..d7f715cdb 100644 --- a/backend/app/modelling/property_filters.py +++ b/backend/app/modelling/property_filters.py @@ -92,17 +92,20 @@ def _resolved_property_types( epcs = session.exec( select(EpcPropertyModel) .where(col(EpcPropertyModel.property_id).in_(property_ids)) - .where(EpcPropertyModel.source == "lodged") + .where(col(EpcPropertyModel.source).in_(("lodged", "predicted"))) ).all() - by_lodged_epc = { - epc.property_id: epc.property_type - for epc in epcs - if epc.property_id is not None and epc.property_type is not None - } + by_epc_source: dict[str, dict[int, str]] = {"lodged": {}, "predicted": {}} + for epc in epcs: + if epc.property_id is not None and epc.property_type is not None: + by_epc_source[epc.source][epc.property_id] = epc.property_type resolved: dict[int, str] = {} for pid in property_ids: - value = by_override.get(pid) or by_lodged_epc.get(pid) + value = ( + by_override.get(pid) + or by_epc_source["lodged"].get(pid) + or by_epc_source["predicted"].get(pid) + ) if value is not None: resolved[pid] = value return resolved