from pip_audit_extra.vulnerability.cache.type import VulnerabilityData, Record

from typing import Any, Set
from datetime import datetime


class JSONObjectHook:
	VULNERABILITY_DATA_FIELDS: Set[str] = {"id", "fix_versions", "severity"}
	RECORD_FIELDS: Set[str] = {"updated_at", "data"}

	def __call__(self, data: Any) -> Any:
		if isinstance(data, dict):
			key_set = set(data.keys())

			if self.VULNERABILITY_DATA_FIELDS == key_set:
				return VulnerabilityData(**data)
			elif self.RECORD_FIELDS == key_set:
				data["updated_at"] = datetime.fromisoformat(data["updated_at"])
				return Record(**data)

		return data
