finalize type hints in helpers

3 files changed, 24 insertions, 30 deletions

diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
index 295a5d84..cbd64a58 100644
--- a/sec_certs/helpers.py
+++ b/sec_certs/helpers.py
@@ -9,7 +9,7 @@ from datetime import date, datetime
 from enum import Enum
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
-from typing import Dict, Hashable, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Dict, Generator, Hashable, Iterator, List, Optional, Sequence, Set, Tuple, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -79,9 +79,9 @@ def get_first_16_bytes_sha256(string: str) -> str:
     return hashlib.sha256(string.encode("utf-8")).hexdigest()[:16]
 
 
-def get_sha256_filepath(filepath):
+def get_sha256_filepath(filepath: Union[str, Path]) -> str:
     hash_sha256 = hashlib.sha256()
-    with open(filepath, "rb") as f:
+    with Path(filepath).open("rb") as f:
         for chunk in iter(lambda: f.read(4096), b""):
             hash_sha256.update(chunk)
     return hash_sha256.hexdigest()
@@ -128,7 +128,6 @@ def sanitize_protection_profiles(record: str) -> list:
     return record.split(",")
 
 
-# TODO: realize whether this stays or goes somewhere else
 def parse_list_of_tables(txt: str) -> Set[str]:
     """
     Parses list of tables from function find_tables(), finds ones that mention algorithms
@@ -162,7 +161,7 @@ def find_tables_iterative(file_text: str) -> List[int]:
     return list(pages)
 
 
-def find_tables(txt: str, file_name: Path) -> Optional[List]:
+def find_tables(txt: str, file_name: Path) -> Optional[Union[List[str], List[int]]]:
     """
     Function that tries to pages in security policy pdf files, where it's possible to find a table containing
     algorithms
@@ -186,7 +185,7 @@ def find_tables(txt: str, file_name: Path) -> Optional[List]:
     return table_page_indices if table_page_indices else None
 
 
-def repair_pdf(file: Path):
+def repair_pdf(file: Path) -> None:
     """
     Some pdfs can't be opened by PyPDF2 - opening them with pikepdf and then saving them fixes this issue.
     By opening this file in a pdf reader, we can already extract number of pages
@@ -197,7 +196,7 @@ def repair_pdf(file: Path):
     pdf.save(file)
 
 
-def convert_pdf_file(pdf_path: Path, txt_path: Path, options):
+def convert_pdf_file(pdf_path: Path, txt_path: Path, options) -> str:
     response = subprocess.run(
         ["pdftotext", *options, pdf_path, txt_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=60
     ).returncode
@@ -207,7 +206,7 @@ def convert_pdf_file(pdf_path: Path, txt_path: Path, options):
         return constants.RETURNCODE_NOK
 
 
-def extract_pdf_metadata(filepath: Path):
+def extract_pdf_metadata(filepath: Path) -> Tuple[str, Optional[Dict[str, Any]]]:
     metadata = dict()
 
     try:
@@ -625,7 +624,7 @@ def plot_dataframe_graph(
     bins: int = 50,
     log: bool = True,
     show: bool = True,
-):
+) -> None:
     pd_data = pd.Series(data)
     pd_data.hist(bins=bins, label=label, density=density, cumulative=cumulative)
     plt.savefig(file_name)
@@ -638,7 +637,7 @@ def plot_dataframe_graph(
     logger.info(sorted_data.where(sorted_data > 1).dropna())
 
 
-def is_in_dict(target_dict, path):
+def is_in_dict(target_dict: Dict, path: str) -> bool:
     current_level = target_dict
     for item in path:
         if item not in current_level:
@@ -648,16 +647,16 @@ def is_in_dict(target_dict, path):
     return True
 
 
-def search_files(folder):
-    for root, dirs, files in os.walk(folder):
+def search_files(folder: str) -> Iterator[str]:
+    for root, _, files in os.walk(folder):
         yield from [os.path.join(root, x) for x in files]
 
 
-def save_modified_cert_file(target_file, modified_cert_file_text, is_unicode_text):
+def save_modified_cert_file(target_file: Union[str, Path], modified_cert_file_text: str, is_unicode_text: bool) -> None:
     if is_unicode_text:
-        write_file = open(target_file, "w", encoding="utf8", errors="replace")
+        write_file = Path(target_file).open("w", encoding="utf8", errors="replace")
     else:
-        write_file = open(target_file, "w", errors="replace")
+        write_file = Path(target_file).open("w", errors="replace")
 
     try:
         write_file.write(modified_cert_file_text)
@@ -739,15 +738,16 @@ def parse_cert_file(file_name, search_rules, limit_max_lines=-1, line_separator=
     return items_found_all, (whole_text_with_newlines, was_unicode_decode_error)
 
 
-def normalize_match_string(match):
+def normalize_match_string(match: str) -> str:
     match = match.strip().rstrip('];.”":)(,').rstrip(os.sep).replace("  ", " ")
     return "".join(filter(str.isprintable, match))
 
 
-def load_cert_file(file_name, limit_max_lines=-1, line_separator=LINE_SEPARATOR):
+def load_cert_file(
+    file_name: Union[str, Path], limit_max_lines: int = -1, line_separator: str = LINE_SEPARATOR
+) -> Tuple[str, str, bool]:
     lines = []
-    was_unicode_decode_error = False
-    with open(file_name, "r", errors=FILE_ERRORS_STRATEGY) as f:
+    with Path(file_name).open("r", errors=FILE_ERRORS_STRATEGY) as f:
         try:
             lines = f.readlines()
         except UnicodeDecodeError:
@@ -786,7 +786,7 @@ def load_cert_file(file_name, limit_max_lines=-1, line_separator=LINE_SEPARATOR)
     return whole_text, whole_text_with_newlines, was_unicode_decode_error
 
 
-def load_cert_html_file(file_name):
+def load_cert_html_file(file_name: str) -> str:
     with open(file_name, "r", errors=FILE_ERRORS_STRATEGY) as f:
         try:
             whole_text = f.read()
@@ -800,7 +800,7 @@ def load_cert_html_file(file_name):
     return whole_text
 
 
-def gen_dict_extract(dct: Dict, searched_key: Hashable = "count"):
+def gen_dict_extract(dct: Dict, searched_key: Hashable = "count") -> Generator[Any, None, None]:
     """
     Function to flatten dictionary with some serious limitations. We only expect to use it temporarily on dictionary
     produced by extract_keywords that contains many layers. On the deepest level in that dictionary, 'some_match': {'count': frequency}.
@@ -857,9 +857,3 @@ def tokenize_dataset(dset: List[str], keywords: Set[str]) -> np.ndarray:
 
 def tokenize(string: str, keywords: Set[str]) -> str:
     return " ".join([x for x in string.split() if x.lower() in keywords])
-
-
-def filter_shortly_described_cves(x, y):
-    n_tokens = np.array(list(map(lambda item: len(item.split(" ")), x)))
-    indices = np.where(n_tokens > 5)
-    return np.array(x)[indices], np.array(y)[indices]
diff --git a/sec_certs/parallel_processing.py b/sec_certs/parallel_processing.py
index 7a150070..6e1c28a6 100644
--- a/sec_certs/parallel_processing.py
+++ b/sec_certs/parallel_processing.py
@@ -1,6 +1,6 @@
 import time
 from multiprocessing.pool import ThreadPool
-from typing import Callable, Iterable, Optional, Union
+from typing import Any, Callable, Iterable, List, Optional, Union
 
 from billiard.pool import Pool
 
@@ -16,7 +16,7 @@ def process_parallel(
     progress_bar: bool = True,
     unpack: bool = False,
     progress_bar_desc: Optional[str] = None,
-):
+) -> List[Any]:
 
     pool: Union[Pool, ThreadPool] = ThreadPool(max_workers) if use_threading else Pool(max_workers)
     results = (
diff --git a/sec_certs/sample/fips.py b/sec_certs/sample/fips.py
index 3bb94b76..3f980fc0 100644
--- a/sec_certs/sample/fips.py
+++ b/sec_certs/sample/fips.py
@@ -479,7 +479,7 @@ class FIPSCertificate(Certificate["FIPSCertificate"], ComplexSerializableType):
             items_found = FIPSCertificate.initialize_dictionary()
             items_found["cert_id"] = int(file.stem)
 
-        text = helpers.load_cert_html_file(file)
+        text = helpers.load_cert_html_file(str(file))
         soup = BeautifulSoup(text, "html.parser")
         for div in soup.find_all("div", class_="row padrow"):
             FIPSCertificate.parse_html_main(div, items_found, pairs)