format: Specified type-hints

2 files changed, 25 insertions, 19 deletions

diff --git a/sec_certs/certificate/common_criteria.py b/sec_certs/certificate/common_criteria.py
index fd97137b..af7688b1 100644
--- a/sec_certs/certificate/common_criteria.py
+++ b/sec_certs/certificate/common_criteria.py
@@ -193,10 +193,10 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
         related_cves: Optional[List[CVE]] = field(default=None)
         cert_lab: Optional[List[str]] = field(default=None)
         cert_id: Optional[str] = field(default=None)
-        directly_affected_by: Optional[List] = field(default=None)
-        indirectly_affected_by: Optional[Set] = field(default=None)
-        directly_affecting: Optional[Set] = field(default=None)
-        indirectly_affecting: Optional[Set] = field(default=None)
+        directly_affected_by: Optional[List[str]] = field(default=None)
+        indirectly_affected_by: Optional[Set[str]] = field(default=None)
+        directly_affecting: Optional[Set[str]] = field(default=None)
+        indirectly_affecting: Optional[Set[str]] = field(default=None)
 
         # manufacturer_list: Optional[List[str]]
 
diff --git a/sec_certs/model/dependency_finder.py b/sec_certs/model/dependency_finder.py
index 24d3bf8d..19a2fc31 100644
--- a/sec_certs/model/dependency_finder.py
+++ b/sec_certs/model/dependency_finder.py
@@ -1,21 +1,27 @@
-from typing import List, Set, Dict, Tuple
+from typing import List, Set, Dict, Tuple, Union, Optional
+from sec_certs.certificate.common_criteria import CommonCriteriaCert
+
+Certificates = Dict[str, CommonCriteriaCert]
+ReferencedByDirect = Dict[str, List[str]]
+ReferencedByIndirect = Dict[str, Set[str]]
+Dependencies = Dict[str, Dict[str, Union[Optional[List[str]], Optional[Set[str]]]]]
 
 
 class DependencyFinder:
 
     def __init__(self):
-        self.dependencies = {}
+        self.dependencies: Dependencies = {}
 
     @staticmethod
-    def _update_direct_references(referenced_by: Dict, cert_id: str, this_cert_id: str) -> None:
+    def _update_direct_references(referenced_by: ReferencedByDirect, cert_id: str, this_cert_id: str) -> None:
         if cert_id not in referenced_by:
             referenced_by[cert_id] = []
         if this_cert_id not in referenced_by[cert_id]:
             referenced_by[cert_id].append(this_cert_id)
 
     @staticmethod
-    def _build_cert_references(certificates: Dict) -> Tuple[Dict, Dict]:
-        referenced_by = {}
+    def _build_cert_references(certificates: Certificates) -> Tuple[ReferencedByDirect, ReferencedByIndirect]:
+        referenced_by: ReferencedByDirect = {}
 
         for cert_obj in certificates.values():
             if cert_obj.pdf_data.report_keywords is None:
@@ -30,7 +36,7 @@ class DependencyFinder:
                 if cert_id != this_cert_id and this_cert_id is not None:
                     DependencyFinder._update_direct_references(referenced_by, cert_id, this_cert_id)
 
-        referenced_by_indirect = {}
+        referenced_by_indirect: ReferencedByIndirect = {}
 
         for cert_id in referenced_by.keys():
             referenced_by_indirect[cert_id] = set()
@@ -55,7 +61,7 @@ class DependencyFinder:
         return referenced_by, referenced_by_indirect
 
     @staticmethod
-    def _get_affecting_directly(cert: str, referenced_by_direct: Dict) -> Set:
+    def _get_affecting_directly(cert: str, referenced_by_direct: ReferencedByDirect) -> Set[str]:
         filter_direct = set()
 
         for cert_id in referenced_by_direct:
@@ -65,7 +71,7 @@ class DependencyFinder:
         return filter_direct
 
     @staticmethod
-    def _get_affecting_indirectly(cert: str, referenced_by_indirect: Dict) -> Set:
+    def _get_affecting_indirectly(cert: str, referenced_by_indirect: ReferencedByIndirect) -> Set[str]:
         filter_indirect = set()
 
         for cert_id in referenced_by_indirect:
@@ -75,14 +81,14 @@ class DependencyFinder:
         return filter_indirect
 
     @staticmethod
-    def _get_affected_directly(cert: str, referenced_by_direct: Dict) -> List:
+    def _get_affected_directly(cert: str, referenced_by_direct: ReferencedByDirect) -> Optional[List[str]]:
         return referenced_by_direct.get(cert, None)
 
     @staticmethod
-    def _get_affected_indirectly(cert: str, referenced_by_indirect: Dict) -> Set:
+    def _get_affected_indirectly(cert: str, referenced_by_indirect: ReferencedByIndirect) -> Optional[Set[str]]:
         return referenced_by_indirect.get(cert, None)
 
-    def fit(self, certificates: Dict) -> None:
+    def fit(self, certificates: Certificates) -> None:
         referenced_by_direct, referenced_by_indirect = DependencyFinder._build_cert_references(certificates)
 
         for dgst in certificates:
@@ -104,14 +110,14 @@ class DependencyFinder:
             self.dependencies[dgst]["indirectly_affecting"] = \
                 DependencyFinder._get_affecting_indirectly(cert_id, referenced_by_indirect)
 
-    def get_directly_affected_by(self, dgst: str) -> List:
+    def get_directly_affected_by(self, dgst: str) -> Optional[List[str]]:
         return self.dependencies[dgst].get("directly_affected_by", None)
 
-    def get_indirectly_affected_by(self, dgst: str) -> Set:
+    def get_indirectly_affected_by(self, dgst: str) -> Optional[Set[str]]:
         return self.dependencies[dgst].get("indirectly_affected_by", None)
 
-    def get_directly_affecting(self, dgst: str) -> Set:
+    def get_directly_affecting(self, dgst: str) -> Optional[Set[str]]:
         return self.dependencies[dgst].get("directly_affecting", None)
 
-    def get_indirectly_affecting(self, dgst: str) -> Set:
+    def get_indirectly_affecting(self, dgst: str) -> Optional[Set[str]]:
         return self.dependencies[dgst].get("indirectly_affecting", None)