make handle_binary_dependencies class method

protectai · Jan 17, 2024 · 1e78ba4 · 1e78ba4
1 parent cadda0f
commit 1e78ba4
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 23 deletions.
diff --git a/modelscan/scanners/h5/scan.py b/modelscan/scanners/h5/scan.py
@@ -77,21 +77,20 @@ def _get_keras_h5_operator_names(self, source: Union[str, Path]) -> List[str]:
 
         return []
 
-    @staticmethod
-    def name() -> str:
-        return "hdf5"
-
-    @staticmethod
-    def full_name() -> str:
-        return "modelscan.scanners.H5LambdaDetectScan"
-
-    @staticmethod
     def handle_binary_dependencies(
-        settings: Optional[Dict[str, Any]] = None
+        self, settings: Optional[Dict[str, Any]] = None
     ) -> Optional[ModelScanError]:
         if not h5py_installed:
             return ModelScanError(
                 H5LambdaDetectScan.name(),
                 f"To use {H5LambdaDetectScan.full_name()}, please install modelscan with h5py extras. 'pip install \"modelscan\[h5py]\"' if you are using pip.",
             )
         return None
+
+    @staticmethod
+    def name() -> str:
+        return "hdf5"
+
+    @staticmethod
+    def full_name() -> str:
+        return "modelscan.scanners.H5LambdaDetectScan"
diff --git a/modelscan/scanners/saved_model/scan.py b/modelscan/scanners/saved_model/scan.py
@@ -87,6 +87,16 @@ def _check_for_unsafe_tf_keras_operator(
             )
         return ScanResults(issues, [])
 
+    def handle_binary_dependencies(
+        self, settings: Optional[Dict[str, Any]] = None
+    ) -> Optional[ModelScanError]:
+        if not tensorflow_installed:
+            return ModelScanError(
+                self.name(),
+                f"To use {self.full_name()}, please install modelscan with tensorflow extras. 'pip install \"modelscan\[tensorflow]\"' if you are using pip.",
+            )
+        return None
+
     @staticmethod
     def name() -> str:
         return "saved_model"
@@ -95,17 +105,6 @@ def name() -> str:
     def full_name() -> str:
         return "modelscan.scanners.SavedModelScan"
 
-    @staticmethod
-    def handle_binary_dependencies(
-        settings: Optional[Dict[str, Any]] = None
-    ) -> Optional[ModelScanError]:
-        if not tensorflow_installed:
-            return ModelScanError(
-                SavedModelScan.name(),
-                f"To use {SavedModelScan.full_name()}, please install modelscan with tensorflow extras. 'pip install \"modelscan\[tensorflow]\"' if you are using pip.",
-            )
-        return None
-
 
 class SavedModelLambdaDetectScan(SavedModelScan):
     def _scan(self, source: Union[str, Path], data: IO[bytes]) -> Optional[ScanResults]:

diff --git a/modelscan/scanners/scan.py b/modelscan/scanners/scan.py
@@ -40,9 +40,8 @@ def scan(
     ) -> Optional[ScanResults]:
         raise NotImplementedError
 
-    @staticmethod
     def handle_binary_dependencies(
-        settings: Optional[Dict[str, Any]] = None
+        self, settings: Optional[Dict[str, Any]] = None
     ) -> Optional[ModelScanError]:
         """
         Implement this method if the plugin requires a binary dependency.