from pip_audit_extra.vulnerability.filter.checker import Checker
from pip_audit_extra.vulnerability.filter.severity import SeverityFilterOption, SeverityChecker, SeverityExacChecker
from pip_audit_extra.vulnerability.dataclass import Vulnerability

from typing import List, Optional, Iterable, Generator


class VulnerabilityFilter:
	checkers: List[Checker]

	def __init__(self, severity: Optional[SeverityFilterOption] = None) -> None:
		self.checkers = []

		if severity is not None:
			if severity.exac:
				self.checkers.append(SeverityExacChecker(severity.value))
			else:
				self.checkers.append(SeverityChecker(severity.value))

	def filter(self, vulns: Iterable[Vulnerability]) -> Generator[Vulnerability, None, None]:
		for vuln in vulns:
			if all(map(lambda x: x.check(vuln), self.checkers)):
				yield vuln
