From 207dac066d0f799bb129c7fa6a4c5bacfc7a9490 Mon Sep 17 00:00:00 2001
From: J08nY
Date: Thu, 27 Apr 2023 12:56:07 +0200
Subject: Compute scheme data matches in heuristics, after cert_ids are
 computed.

---
 src/sec_certs/dataset/cc.py | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

(limited to 'src')

diff --git a/src/sec_certs/dataset/cc.py b/src/sec_certs/dataset/cc.py
index 161c8f65..0b2e1ae2 100644
--- a/src/sec_certs/dataset/cc.py
+++ b/src/sec_certs/dataset/cc.py
@@ -698,9 +698,23 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
             self.certs[dgst].heuristics.direct_transitive_cves = transitive_cve.direct_transitive_cves
             self.certs[dgst].heuristics.indirect_transitive_cves = transitive_cve.indirect_transitive_cves
 
+    def _compute_scheme_data(self):
+        for scheme in self.auxiliary_datasets.scheme_dset:
+            if certified := scheme.lists.get(EntryType.Certified):
+                certs = [cert for cert in self if cert.status == "active"]
+                matches = CCSchemeMatcher.match_all(certified, scheme.country, certs)
+                for dgst, match in matches.items():
+                    self[dgst].heuristics.scheme_data = match
+            if archived := scheme.lists.get(EntryType.Archived):
+                certs = [cert for cert in self if cert.status == "archived"]
+                matches = CCSchemeMatcher.match_all(archived, scheme.country, certs)
+                for dgst, match in matches.items():
+                    self[dgst].heuristics.scheme_data = match
+
     def _compute_heuristics(self) -> None:
         self._compute_normalized_cert_ids()
         super()._compute_heuristics()
+        self._compute_scheme_data()
         self._compute_cert_labs()
         self._compute_sars()
 
@@ -823,18 +837,6 @@ class CCDataset(Dataset[CCCertificate, CCAuxiliaryDatasets], ComplexSerializable
         else:
             scheme_dset = CCSchemeDataset.from_json(self.scheme_dataset_path)
 
-        for scheme in scheme_dset:
-            if certified := scheme.lists.get(EntryType.Certified):
-                certs = [cert for cert in self if cert.status == "active"]
-                matches = CCSchemeMatcher.match_all(certified, scheme.country, certs)
-                for dgst, match in matches.items():
-                    self[dgst].heuristics.scheme_data = match
-            if archived := scheme.lists.get(EntryType.Archived):
-                certs = [cert for cert in self if cert.status == "archived"]
-                matches = CCSchemeMatcher.match_all(archived, scheme.country, certs)
-                for dgst, match in matches.items():
-                    self[dgst].heuristics.scheme_data = match
-
         return scheme_dset
 
 
-- 
cgit v1.3.1