Refactor importer/improver registry

Signed-off-by: Hritik Vijay <hritikxx8@gmail.com>

Author: Hritik14

vulnerabilities/import_runner.py

@@ -25,7 +25,7 @@
 import datetime
 import json
 import logging
-from typing import Set
+from typing import List
 from typing import Iterable
 
 
@@ -55,18 +55,22 @@ def run(self) -> None:
         """
         Create a data source for the given importer and store the data retrieved in the database.
         """
-        logger.info(f"Starting import for {self.importer.qualified_name}")
-        advisory_datas = self.importer().advisory_data()
         importer_name = self.importer.qualified_name
-        process_advisories(advisory_datas=advisory_datas, importer_name=importer_name)
-        logger.info(f"Finished import for {self.importer.qualified_name}.")
+        importer_class = self.importer
+        logger.info(f"Starting import for {importer_name}")
+        advisory_datas = importer_class().advisory_data()
+        count = process_advisories(advisory_datas=advisory_datas, importer_name=importer_name)
+        logger.info(
+            f"Finished import for {importer_name}. Imported {count} advisories."
+        )
 
 
-def process_advisories(advisory_datas: Iterable[AdvisoryData], importer_name: str) -> None:
+def process_advisories(advisory_datas: Iterable[AdvisoryData], importer_name: str) -> List:
     """
     Insert advisories into the database
+    Return list of ids of inserted advisories
     """
-
+    count = 0
     for data in advisory_datas:
         obj, created = Advisory.objects.get_or_create(
             aliases=data.aliases,
@@ -83,5 +87,8 @@ def process_advisories(advisory_datas: Iterable[AdvisoryData], importer_name: st
             logger.info(
                 f"[*] New Advisory with aliases: {obj.aliases!r}, created_by: {obj.created_by}"
             )
+            count += 1
         else:
             logger.debug(f"Advisory with aliases: {obj.aliases!r} already exists. Skipped.")
+
+    return count

vulnerabilities/importers/__init__.py

@@ -21,6 +21,6 @@
 #  Visit https://github.com/nexB/vulnerablecode/ for support and download.
 from vulnerabilities.importers import nginx
 
-IMPORTER_REGISTRY = [nginx.NginxDataSource]
+IMPORTERS = [nginx.NginxDataSource]
 
-importer_mapping = {x.qualified_name: x for x in IMPORTER_REGISTRY}
+IMPORTER_REGISTRY = {x.qualified_name: x for x in IMPORTERS}

vulnerabilities/improvers/__init__.py

@@ -1,6 +1,6 @@
 from vulnerabilities.improvers import default
 from vulnerabilities import importers
 
-IMPROVER_REGISTRY = [default.DefaultImprover, importers.nginx.NginxBasicImprover]
+IMPROVERS = [default.DefaultImprover, importers.nginx.NginxBasicImprover]
 
-improver_mapping = {x.qualified_name: x for x in IMPROVER_REGISTRY}
+IMPROVER_REGISTRY = {x.qualified_name: x for x in IMPROVERS}

vulnerabilities/management/commands/import.py

@@ -27,7 +27,6 @@
 from django.core.management.base import CommandError
 
 from vulnerabilities.importers import IMPORTER_REGISTRY
-from vulnerabilities.importers import importer_mapping
 from vulnerabilities.import_runner import ImportRunner
 
 
@@ -54,7 +53,7 @@ def handle(self, *args, **options):
             return
 
         if options["all"]:
-            self.import_data(IMPORTER_REGISTRY)
+            self.import_data(IMPORTER_REGISTRY.values())
             return
 
         sources = options["sources"]
@@ -63,14 +62,18 @@ def handle(self, *args, **options):
                 'Please provide at least one data source to import from or use "--all".'
             )
 
-        self.import_data(valid_sources(sources))
+        self.import_data(validate_importers(sources))
 
     def list_sources(self):
-        importers = [importer.qualified_name for importer in IMPORTER_REGISTRY]
+        importers = list(IMPORTER_REGISTRY)
         self.stdout.write("Vulnerability data can be imported from the following sources:")
         self.stdout.write("\n".join(importers))
 
     def import_data(self, importers):
+        """
+        Run the given ``importers``. The ``importers`` are expected to be class
+        names for the importers.
+        """
         failed_importers = []
 
         for importer in importers:
@@ -95,12 +98,12 @@ def import_data(self, importers):
             raise CommandError(f"{len(failed_importers)} failed!: {','.join(failed_importers)}")
 
 
-def valid_sources(sources):
+def validate_importers(sources):
     importers = []
     unknown_sources = []
     for source in sources:
         try:
-            importers.append(importer_mapping[source])
+            importers.append(IMPORTER_REGISTRY[source])
         except KeyError:
             unknown_sources.append(source)
     if unknown_sources:

vulnerabilities/management/commands/improve.py

@@ -28,7 +28,6 @@
 
 from vulnerabilities.import_runner import ImportRunner
 from vulnerabilities.improvers import IMPROVER_REGISTRY
-from vulnerabilities.improvers import improver_mapping
 from vulnerabilities.improve_runner import ImproveRunner
 
 
@@ -52,17 +51,17 @@ def handle(self, *args, **options):
             return
 
         if options["all"]:
-            self.improve_data(IMPROVER_REGISTRY)
+            self.improve_data(IMPROVER_REGISTRY.values())
             return
 
         sources = options["sources"]
         if not sources:
             raise CommandError('Please provide at least one improver to run use "--all".')
 
-        self.improve_data(valid_sources(sources))
+        self.improve_data(validate_improvers(sources))
 
     def list_sources(self):
-        improvers = [improver.qualified_name for improver in IMPROVER_REGISTRY]
+        improvers = list(IMPROVER_REGISTRY)
         self.stdout.write("Vulnerability data can be processed by these available improvers:\n")
         self.stdout.write("\n".join(improvers))
 
@@ -91,12 +90,12 @@ def improve_data(self, improvers):
             raise CommandError(f"{len(failed_improvers)} failed!: {','.join(failed_improvers)}")
 
 
-def valid_sources(sources):
+def validate_improvers(sources):
     improvers = []
     unknown_sources = []
     for source in sources:
         try:
-            improvers.append(improver_mapping[source])
+            improvers.append(IMPROVER_REGISTRY[source])
         except KeyError:
             unknown_sources.append(source)
     if unknown_sources: