fix notebooks

5 files changed, 22 insertions, 27 deletions

diff --git a/notebooks/cc/cpe_eval.ipynb b/notebooks/cc/cpe_eval.ipynb
index 088d4659..de948a45 100644
--- a/notebooks/cc/cpe_eval.ipynb
+++ b/notebooks/cc/cpe_eval.ipynb
@@ -18,7 +18,8 @@
     "from sec_certs.dataset import CCDataset\n",
     "import pandas as pd\n",
     "import json\n",
-    "import tempfile"
+    "import tempfile\n",
+    "from sec_certs.utils.label_studio_utils import to_label_studio_json"
    ]
   },
   {
@@ -58,7 +59,7 @@
     "with tempfile.TemporaryDirectory() as tmp_dir:\n",
     "    dset.root_dir = tmp_dir\n",
     "    dset.certs = {x.dgst: x for x in dset if x.dgst in eval_certs.index.tolist()}\n",
-    "    dset.to_label_studio_json(\"./label_studio_input_data.json\", update_json=False)"
+    "    to_label_studio_json(dset, \"./label_studio_input_data.json\")"
    ]
   },
   {
diff --git a/notebooks/cc/scheme_eval.ipynb b/notebooks/cc/scheme_eval.ipynb
index 7ffc6f16..9150cb8c 100644
--- a/notebooks/cc/scheme_eval.ipynb
+++ b/notebooks/cc/scheme_eval.ipynb
@@ -24,7 +24,8 @@
     "from sec_certs.model import CCSchemeMatcher\n",
     "from sec_certs.sample.cc_certificate_id import canonicalize\n",
     "from sec_certs.sample.cc_scheme import CCScheme, EntryType\n",
-    "from sec_certs.configuration import config"
+    "from sec_certs.configuration import config\n",
+    "from sec_certs.dataset.auxiliary_dataset_handling import CCSchemeDatasetHandler"
    ]
   },
   {
@@ -56,7 +57,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dset.auxiliary_datasets.scheme_dset = schemes\n",
+    "dset.aux_handlers[CCSchemeDatasetHandler].dset = schemes\n",
     "\n",
     "count_was = 0\n",
     "count_is = 0\n",
@@ -161,7 +162,7 @@
     "        rate = len(assigned)/len(total) * 100 if len(total) != 0 else 0\n",
     "        rate_list = rates.setdefault(country, [])\n",
     "        rate_list.append(rate)\n",
-    "        \n",
+    "\n",
     "        print(f\"{country}: {len(assigned)} assigned out of {len(total)} -> {rate:.1f}%\")\n",
     "        total_active = total[total[\"status\"] == \"active\"]\n",
     "        assigned_active = assigned[assigned[\"status\"] == \"active\"]\n",
diff --git a/notebooks/cc/vulnerabilities.ipynb b/notebooks/cc/vulnerabilities.ipynb
index 176b3e69..c5a4bfb8 100644
--- a/notebooks/cc/vulnerabilities.ipynb
+++ b/notebooks/cc/vulnerabilities.ipynb
@@ -33,6 +33,7 @@
     "import warnings\n",
     "from pathlib import Path\n",
     "import tempfile\n",
+    "from sec_certs.dataset.auxiliary_dataset_handling import CVEDatasetHandler, CPEDatasetHandler, CCMaintenanceUpdateDatasetHandler\n",
     "from sec_certs.dataset import CCDataset, CCDatasetMaintenanceUpdates, CVEDataset, CPEDataset\n",
     "from sec_certs.utils.pandas import (\n",
     "    compute_cve_correlations,\n",
@@ -82,15 +83,11 @@
     "\n",
     "# # Remote instantiation (takes approx. 10 minutes to complete)\n",
     "# dset: CCDataset = CCDataset.from_web_latest(path=\"dset\", auxiliary_datasets=True)\n",
+    "# dset.load_auxiliary_datasets()\n",
     "\n",
-    "# print(\"Downloading dataset of maintenance updates\")\n",
-    "# main_dset: CCDatasetMaintenanceUpdates = CCDatasetMaintenanceUpdates.from_web_latest()\n",
-    "\n",
-    "# print(\"Downloading CPE dataset\")\n",
-    "# cpe_dset: CPEDataset = dset.auxiliary_datasets.cpe_dset\n",
-    "\n",
-    "# print(\"Downloading CVE dataset\")\n",
-    "# cve_dset: CVEDataset = dset.auxiliary_datasets.cve_dset"
+    "# main_dset: CCDatasetMaintenanceUpdates = dset.aux_handlers[CCMaintenanceUpdateDatasetHandler].dset\n",
+    "# cpe_dset: CPEDataset = dset.aux_handlers[CPEDatasetHandler].dset\n",
+    "# cve_dset: CVEDataset = dset.aux_handlers[CVEDatasetHandler].dset"
    ]
   },
   {
diff --git a/notebooks/fips/vulnerabilities.ipynb b/notebooks/fips/vulnerabilities.ipynb
index 3b3d21e3..d495ea6a 100644
--- a/notebooks/fips/vulnerabilities.ipynb
+++ b/notebooks/fips/vulnerabilities.ipynb
@@ -1,10 +1,12 @@
 {
  "cells": [
   {
-   "metadata": {},
    "cell_type": "markdown",
-   "source": "# Vulnerability analysis",
-   "id": "3a0981d008383c12"
+   "id": "3a0981d008383c12",
+   "metadata": {},
+   "source": [
+    "# Vulnerability analysis"
+   ]
   },
   {
    "cell_type": "code",
@@ -21,6 +23,7 @@
     "from sec_certs.dataset.fips import FIPSDataset\n",
     "from sec_certs.dataset.cpe import CPEDataset\n",
     "from sec_certs.dataset.cve import CVEDataset\n",
+    "from sec_certs.dataset.auxiliary_dataset_handling import CPEDatasetHandler, CVEDatasetHandler\n",
     "from sec_certs.utils.pandas import expand_df_with_cve_cols\n",
     "import pandas as pd\n",
     "import seaborn as sns\n",
@@ -47,8 +50,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cve_dset: CVEDataset = dset.auxiliary_datasets.cve_dset\n",
-    "cpe_dset: CPEDataset = dset.auxiliary_datasets.cpe_dset"
+    "dset.load_auxiliary_datasets()\n",
+    "cve_dset: CVEDataset = dset.aux_handlers[CVEDatasetHandler].dset\n",
+    "cpe_dset: CPEDataset = dset.aux_handlers[CPEDatasetHandler].dset"
    ]
   },
   {
@@ -181,14 +185,6 @@
     "g = sns.relplot(data=df_cve_rich, x=\"level\", y=\"avg_cve_score\")\n",
     "plt.show()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6c3c2ec4-3fab-48ad-aacb-6f54277abe66",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
diff --git a/src/sec_certs/dataset/cc.py b/src/sec_certs/dataset/cc.py
index 529beb85..4d09106d 100644
--- a/src/sec_certs/dataset/cc.py
+++ b/src/sec_certs/dataset/cc.py
@@ -845,7 +845,7 @@ class CCDataset(Dataset[CCCertificate], ComplexSerializableType):
         compute_transitive_vulnerabilities(self.certs)
 
         if not skip_schemes:
-            compute_scheme_data(self.aux_handlers[CCSchemeDatasetHandler].dset, self.certs.values())
+            compute_scheme_data(self.aux_handlers[CCSchemeDatasetHandler].dset, self.certs)
 
         compute_cert_labs(self.certs.values())
         compute_sars(self.certs.values())