replace some ifs with dictionaray setdefault

4 files changed, 6 insertions, 15 deletions

diff --git a/src/sec_certs/dataset/cpe.py b/src/sec_certs/dataset/cpe.py
index 62f4d500..aadac766 100644
--- a/src/sec_certs/dataset/cpe.py
+++ b/src/sec_certs/dataset/cpe.py
@@ -117,11 +117,8 @@ class CPEDataset(JSONPathDataset, ComplexSerializableType):
         return pd.DataFrame([x.pandas_tuple for x in self], columns=CPE.pandas_columns).set_index("uri")
 
     def get_title_to_cpes_dict(self) -> dict[str, set[CPE]]:
-        title_to_cpes_dict = {}
+        title_to_cpes_dict: dict[str, set[CPE]] = {}
         for cpe in self:
             if cpe.title:
-                if cpe.title not in title_to_cpes_dict:
-                    title_to_cpes_dict[cpe.title] = {cpe}
-                else:
-                    title_to_cpes_dict[cpe.title].add(cpe)
+                title_to_cpes_dict.setdefault(cpe.title, set()).add(cpe)
         return title_to_cpes_dict
diff --git a/src/sec_certs/dataset/cve.py b/src/sec_certs/dataset/cve.py
index 4e94a5d9..8bf3b40e 100644
--- a/src/sec_certs/dataset/cve.py
+++ b/src/sec_certs/dataset/cve.py
@@ -138,10 +138,7 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType):
                 continue
 
             for cpe_uri in vulnerable_cpe_uris:
-                if cpe_uri in self._cpe_uri_to_cve_ids_lookup:
-                    self._cpe_uri_to_cve_ids_lookup[cpe_uri].add(cve.cve_id)
-                else:
-                    self._cpe_uri_to_cve_ids_lookup[cpe_uri] = {cve.cve_id}
+                self._cpe_uri_to_cve_ids_lookup.setdefault(cpe_uri, set()).add(cve.cve_id)
 
     def _get_cves_from_exactly_matched_cpes(self, cpe_uris: set[str]) -> set[str]:
         return set(
diff --git a/src/sec_certs/model/cpe_matching.py b/src/sec_certs/model/cpe_matching.py
index 3015aeaa..1a02008b 100644
--- a/src/sec_certs/model/cpe_matching.py
+++ b/src/sec_certs/model/cpe_matching.py
@@ -70,10 +70,7 @@ class CPEClassifier:
 
         for cpe in tqdm(sufficiently_long_cpes, desc="Fitting the CPE classifier"):
             self.vendor_to_versions_[cpe.vendor].add(cpe.version)
-            if (cpe.vendor, cpe.version) not in self.vendor_version_to_cpe_:
-                self.vendor_version_to_cpe_[(cpe.vendor, cpe.version)] = {cpe}
-            else:
-                self.vendor_version_to_cpe_[(cpe.vendor, cpe.version)].add(cpe)
+            self.vendor_version_to_cpe_.setdefault((cpe.vendor, cpe.version), set()).add(cpe)
 
     def predict(self, X: list[tuple[str, str, str]]) -> list[set[str] | None]:
         """
diff --git a/src/sec_certs/serialization/json.py b/src/sec_certs/serialization/json.py
index 7f86ab3b..314352bb 100644
--- a/src/sec_certs/serialization/json.py
+++ b/src/sec_certs/serialization/json.py
@@ -95,7 +95,7 @@ class ComplexSerializableType:
 # Decorator for serialization
 def serialize(func: Callable):
     @wraps(func)
-    def inner_func(*args, **kwargs):
+    def _serialize(*args, **kwargs):
         if not args or not issubclass(type(args[0]), ComplexSerializableType):
             raise ValueError(
                 "@serialize decorator is to be used only on instance methods of ComplexSerializableType child classes."
@@ -112,7 +112,7 @@ def serialize(func: Callable):
             args[0].to_json()
         return result
 
-    return inner_func
+    return _serialize
 
 
 def get_class_fullname(obj: Any) -> str: