WiP new cve and cpe dataset handling

1 files changed, 32 insertions, 8 deletions

diff --git a/src/sec_certs/serialization/json.py b/src/sec_certs/serialization/json.py
index 7b523b9b..7f86ab3b 100644
--- a/src/sec_certs/serialization/json.py
+++ b/src/sec_certs/serialization/json.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import copy
+import gzip
 import json
-from datetime import date
+from datetime import date, datetime
 from functools import wraps
 from pathlib import Path
 from typing import Any, Callable, TypeVar
@@ -44,7 +45,12 @@ class ComplexSerializableType:
         except TypeError as e:
             raise TypeError(f"Dict: {dct} on {cls.__mro__}") from e
 
-    def to_json(self, output_path: str | Path | None = None) -> None:
+    def to_json(self, output_path: str | Path | None = None, compress: bool = False) -> None:
+        """
+        Serializes `ComplexSerializableType` instance to json file.
+        :param str | Path | None output_path: path where the file will be stored. If None, `obj.json_path` access is attempted, defaults to None
+        :param bool compress: if True, will be compress with gzip, defaults to False
+        """
         if not output_path and (not hasattr(self, "json_path") or not self.json_path):  # type: ignore
             raise SerializationError(
                 f"The object {self} of type {self.__class__} does not have json_path attribute set but to_json() was called without an argument."
@@ -60,14 +66,30 @@ class ComplexSerializableType:
             raise SerializationError("output path for json cannot be directory.")
 
         # false positive MyPy warning, cannot be None
-        with Path(output_path).open("w") as handle:  # type: ignore
-            json.dump(self, handle, indent=4, cls=CustomJSONEncoder, ensure_ascii=False)
+        if compress:
+            with gzip.open(str(output_path), "w") as handle:  # type: ignore
+                json_str = json.dumps(self, indent=4, cls=CustomJSONEncoder, ensure_ascii=False)
+                handle.write(json_str.encode("utf-8"))
+        else:
+            with Path(output_path).open("w") as handle:  # type: ignore
+                json.dump(self, handle, indent=4, cls=CustomJSONEncoder, ensure_ascii=False)
 
     @classmethod
-    def from_json(cls: type[T], input_path: str | Path) -> T:
-        input_path = Path(input_path)
-        with input_path.open("r") as handle:
-            return json.load(handle, cls=CustomJSONDecoder)
+    def from_json(cls: type[T], input_path: str | Path, is_compressed: bool = False) -> T:
+        """
+        Will load `ComplexSerializableType` from json.
+        :param str | Path input_path: path to load the file from
+        :param bool is_compressed: if True, will decompress .gz first, defaults to False
+        :return T: the deserialized object
+        """
+        if is_compressed:
+            with gzip.open(str(input_path)) as handle:
+                json_str = handle.read().decode("utf-8")
+                return json.loads(json_str, cls=CustomJSONDecoder)
+        else:
+            input_path = Path(input_path)
+            with input_path.open("r") as handle:
+                return json.load(handle, cls=CustomJSONDecoder)
 
 
 # Decorator for serialization
@@ -113,6 +135,8 @@ class CustomJSONEncoder(json.JSONEncoder):
             return sorted(obj)
         if isinstance(obj, date):
             return str(obj)
+        if isinstance(obj, datetime):
+            return obj.isoformat()
         if isinstance(obj, Path):
             return str(obj)
         return super().default(obj)