type hinting checks

4 files changed, 27 insertions, 29 deletions

diff --git a/sec_certs/cert_processing.py b/sec_certs/cert_processing.py
index bc11083b..0ccac952 100644
--- a/sec_certs/cert_processing.py
+++ b/sec_certs/cert_processing.py
@@ -1,9 +1,11 @@
 from tqdm import tqdm
 from multiprocessing.pool import Pool, ThreadPool
+from typing import Callable, Iterable, Optional
 import time
 
-# TODO: Add timeout. Kinda meh with ThreadingTimeout, SignalTimeout does not work on Windows, stopit package.
-def process_parallel(func, items, max_workers, callback=None, use_threading=True, progress_bar=True):
+
+def process_parallel(func: Callable, items: Iterable, max_workers: int, callback: Optional[Callable] = None,
+                     use_threading: bool = True, progress_bar: bool = True):
     if use_threading is True:
         pool = ThreadPool(max_workers)
     else:
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 724a60e0..a9ce04b7 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -382,13 +382,13 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
         def to_dict(self):
             return copy.deepcopy(self.__dict__)
 
-        def __lt__(self, other):
-            return self.pp_name < other.pp_name
-
         @classmethod
         def from_dict(cls, dct):
             return cls(*tuple(dct.values()))
 
+        def __lt__(self, other):
+            return self.pp_name < other.pp_name
+
     def __init__(self, category: str, name: str, manufacturer: str, scheme: str,
                  security_level: Union[str, set], not_valid_before: date,
                  not_valid_after: date, report_link: str, st_link: str, src: str, cert_link: Optional[str],
@@ -450,7 +450,7 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
             self.src = self.src + ' + ' + other.src
 
     @classmethod
-    def from_dict(cls, dct: dict) -> 'CommonCriteriaCert':
+    def from_dict(cls, dct: Dict) -> 'CommonCriteriaCert':
         new_dct = dct.copy()
         new_dct['maintainance_updates'] = set(dct['maintainance_updates'])
         new_dct['protection_profiles'] = set(dct['protection_profiles'])
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 85b95b7c..277a84a9 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -3,7 +3,7 @@ import re
 from datetime import datetime
 import locale
 import logging
-from typing import Dict, List, ClassVar, Collection, TypeVar, Type, Union
+from typing import Dict, List, ClassVar, Collection, TypeVar, Type, Union, Generic, Optional, Sequence
 import json
 from importlib import import_module
 
@@ -17,7 +17,7 @@ import requests
 
 from tabula import read_pdf
 import pandas as pd
-from bs4 import BeautifulSoup
+from bs4 import BeautifulSoup, Tag
 
 
 from sec_certs.files import search_files
@@ -116,7 +116,7 @@ class Dataset(ABC):
                 logger.info(f'Failed to convert {path}, exit code: {e}')
 
     @staticmethod
-    def _download_parallel(urls, paths, prune_corrupted=True):
+    def _download_parallel(urls: Collection[str], paths: Collection[Path], prune_corrupted: bool = True):
         exit_codes = cert_processing.process_parallel(download.download_file,
                                                       list(zip(urls, paths)),
                                                       constants.N_THREADS)
@@ -135,10 +135,6 @@ class Dataset(ABC):
 
 
 class CCDataset(Dataset, ComplexSerializableType):
-    def __init__(self, certs: Dict[str, 'CommonCriteriaCert'], root_dir: Path, name: str = 'dataset name',
-                 description: str = 'dataset_description'):
-        super().__init__(certs, root_dir, name, description)
-
     @property
     def web_dir(self) -> Path:
         return self.root_dir / 'web'
@@ -223,7 +219,8 @@ class CCDataset(Dataset, ComplexSerializableType):
         logger.info(
             f'Added {len(will_be_added)} new and merged further {n_merged} certificates to the dataset.')
 
-    def get_certs_from_web(self, to_download=True, keep_metadata: bool = True, get_active=True, get_archived=True):
+    def get_certs_from_web(self, to_download: bool = True, keep_metadata: bool = True, get_active: bool = True,
+                           get_archived: bool = True):
         """
         Downloads all metadata about certificates from CSV and HTML sources
         """
@@ -264,15 +261,13 @@ class CCDataset(Dataset, ComplexSerializableType):
         if not keep_metadata:
             shutil.rmtree(self.web_dir)
 
-    def _get_all_certs_from_csv(self, get_active, get_archived) -> Dict[str, 'CommonCriteriaCert']:
+    def _get_all_certs_from_csv(self, get_active: bool, get_archived: bool) -> Dict[str, 'CommonCriteriaCert']:
         """
         Creates dictionary of new certificates from csv sources.
         """
         csv_sources = self.csv_products.keys()
-        csv_sources = [
-            x for x in csv_sources if 'active' not in x or get_active]
-        csv_sources = [
-            x for x in csv_sources if 'archived' not in x or get_archived]
+        csv_sources = [x for x in csv_sources if 'active' not in x or get_active]
+        csv_sources = [x for x in csv_sources if 'archived' not in x or get_archived]
 
         new_certs = {}
         for file in csv_sources:
@@ -288,7 +283,7 @@ class CCDataset(Dataset, ComplexSerializableType):
         Using pandas, this parses a single CSV file.
         """
 
-        def _get_primary_key_str(row):
+        def _get_primary_key_str(row: Tag):
             prim_key = row['category'] + row['cert_name'] + row['report_link']
             return prim_key
 
@@ -335,15 +330,13 @@ class CCDataset(Dataset, ComplexSerializableType):
                  df_base.itertuples()}
         return certs
 
-    def _get_all_certs_from_html(self, get_active, get_archived) -> Dict[str, 'CommonCriteriaCert']:
+    def _get_all_certs_from_html(self, get_active: bool, get_archived: bool) -> Dict[str, 'CommonCriteriaCert']:
         """
         Prepares dictionary of certificates from all html files.
         """
         html_sources = self.html_products.keys()
-        html_sources = [
-            x for x in html_sources if 'active' not in x or get_active]
-        html_sources = [
-            x for x in html_sources if 'archived' not in x or get_archived]
+        html_sources = [x for x in html_sources if 'active' not in x or get_active]
+        html_sources = [x for x in html_sources if 'archived' not in x or get_archived]
 
         new_certs = {}
         for file in html_sources:
@@ -422,12 +415,14 @@ class CCDataset(Dataset, ComplexSerializableType):
     def _download_reports(self):
         self.reports_pdf_dir.mkdir(parents=True, exist_ok=True)
         reports_urls = [x.report_link for x in self]
-        self._download_parallel(reports_urls, self.report_pdf_paths.values(), prune_corrupted=True)
+        # for noqa below, see: https://youtrack.jetbrains.com/issue/PY-41771
+        self._download_parallel(reports_urls, self.report_pdf_paths.values(), prune_corrupted=True) # noqa
 
     def _download_targets(self):
         self.targets_pdf_dir.mkdir(parents=True, exist_ok=True)
-        target_urls = [x.st_link for x in self]
-        self._download_parallel(target_urls, self.target_pdf_paths.values(), prune_corrupted=True)
+        target_urls = [x.st_link for x in self.certs]
+        # for noqa below, see: https://youtrack.jetbrains.com/issue/PY-41771
+        self._download_parallel(target_urls, self.target_pdf_paths.values(), prune_corrupted=True) # noqa
 
     def download_all_pdfs(self):
         logger.info('Downloading CC certificate reports')
diff --git a/sec_certs/serialization.py b/sec_certs/serialization.py
index d37dd067..c96fc07c 100644
--- a/sec_certs/serialization.py
+++ b/sec_certs/serialization.py
@@ -1,6 +1,7 @@
 import json
 from datetime import date
 from pathlib import Path
+from typing import Dict
 
 from abc import ABC, abstractmethod
 
@@ -13,7 +14,7 @@ class ComplexSerializableType(ABC):
 
     @classmethod
     @abstractmethod
-    def from_dict(cls):
+    def from_dict(cls, dct: Dict):
         raise NotImplementedError