from pip_audit_extra.vulnerability.cache.type import VulnerabilityData

from json import JSONEncoder as _JSONEncoder
from datetime import datetime
from typing import Any


class JSONEncoder(_JSONEncoder):
	def default(self, o: Any) -> Any:
		if isinstance(o, VulnerabilityData):
			return {
				"id": o.id,
				"fix_versions": o.fix_versions,
				"severity": o.severity,
			}
		elif isinstance(o, datetime):
			return o.isoformat()

		return super().default(o)
