Error Handling Around Scan (#116)

protectai · Mar 19, 2024 · 7536ddc · 7536ddc
1 parent 9eb0604
commit 7536ddc
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/modelscan/modelscan.py b/modelscan/modelscan.py
@@ -122,10 +122,24 @@ def _scan_source(
         scanned = False
         for scan_class in self._scanners_to_run:
             scanner = scan_class(self._settings)  # type: ignore[operator]
-            scan_results = scanner.scan(
-                source=source,
-                data=data,
-            )
+
+            try:
+                scan_results = scanner.scan(
+                    source=source,
+                    data=data,
+                )
+            except Exception as e:
+                logger.error(
+                    f"Error encountered from scanner {scanner.full_name()} with path {source}: {e}"
+                )
+                self._errors.append(
+                    ModelScanError(
+                        scanner.full_name(),
+                        ErrorCategories.MODEL_SCAN,
+                        f"Error encountered from scanner {scanner.full_name()}: {e}",
+                        f"{source}",
+                    )
+                )
 
             if scan_results is not None:
                 scanned = True

diff --git a/modelscan/scanners/pickle/scan.py b/modelscan/scanners/pickle/scan.py
@@ -67,8 +67,11 @@ def scan(
         if data:
             results = scan_numpy(data=data, source=source, settings=self._settings)
 
-        with open(source, "rb") as file_io:
-            results = scan_numpy(data=file_io, source=source, settings=self._settings)
+        else:
+            with open(source, "rb") as file_io:
+                results = scan_numpy(
+                    data=file_io, source=source, settings=self._settings
+                )
 
         return self.label_results(results)