Feat: add vetted properties

.gitignore

@@ -0,0 +1 @@
+__pycache__/

example/vetted.py

@@ -0,0 +1,15 @@
+from fickling.pickle import Pickled
+import pickle
+import numpy as np
+
+
+arr = np.array([1,2,3])
+fickled = Pickled.load(pickle.dumps(arr))
+fickled.vetted_dependencies = ["numpy"]
+fickled.vetted_calls = ["ndarray", "dtype"]
+
+is_safe = fickled.is_likely_safe
+if is_safe:
+    print("✅")
+else:
+    print("❌")

fickling/pickle.py

@@ -246,6 +246,8 @@ def __init__(self, opcodes: Iterable[Opcode]):
         self._opcodes: List[Opcode] = list(opcodes)
         self._ast: Optional[ast.Module] = None
         self._properties: Optional[ASTProperties] = None
+        self._vetted_dependencies: Optional[List[str]] = None
+        self._vetted_calls: Optional[List[str]] = None
 
     def __len__(self) -> int:
         return len(self._opcodes)
@@ -404,11 +406,37 @@ def has_import(self) -> bool:
         """Checks whether unpickling would cause an import to be run"""
         return bool(self.properties.imports)
 
+    @property
+    def has_unvetted_import(self) -> bool:
+        if self._vetted_dependencies is None:
+            raise Exception("Cannot call has_unvetted_import when vetted_dependencies is not set")
+        unvetted_import = False
+        for node in self.properties.imports:
+            module_path = node.module.split(".")
+            if module_path[0] not in self._vetted_dependencies:
+                unvetted_import = True
+                break
+        return unvetted_import
+
     @property
     def has_call(self) -> bool:
         """Checks whether unpickling would cause a function call"""
         return bool(self.properties.calls)
 
+    @property
+    def has_unvetted_calls(self) -> bool:
+        if self._vetted_calls is None:
+            raise Exception("Cannot call has_unvetted_functions when vetted_functions is not set")
+        unvetted_call = False
+        for call in self.properties.non_setstate_calls:
+            function = call.func.id
+            if function == "_reconstruct":
+                function = call.args[0].id
+            if function not in self._vetted_calls:
+                unvetted_call = True
+                break
+        return unvetted_call
+
     @property
     def has_non_setstate_call(self) -> bool:
         """Checks whether unpickling would cause a call to a function other than object.__setstate__"""
@@ -417,7 +445,10 @@ def has_non_setstate_call(self) -> bool:
     @property
     def is_likely_safe(self) -> bool:
         # `self.has_call` is probably safe as long as `not self.has_import`
-        return not self.has_import and not self.has_non_setstate_call
+        if self._vetted_dependencies is None or self._vetted_calls is None:
+            return not self.has_import and not self.has_non_setstate_call
+        else:
+            return not self.has_unvetted_import and not self.has_unvetted_calls
 
     def unsafe_imports(self) -> Iterator[Union[ast.Import, ast.ImportFrom]]:
         for node in self.properties.imports:
@@ -437,6 +468,22 @@ def ast(self) -> ast.Module:
             self._ast = Interpreter.interpret(self)
         return self._ast
 
+    @property
+    def vetted_dependencies(self) -> Optional[List[str]]:
+        return self._vetted_dependencies
+
+    @vetted_dependencies.setter
+    def vetted_dependencies(self, dependencies: List[str]):
+        self._vetted_dependencies = dependencies
+
+    @property
+    def vetted_calls(self) -> Optional[List[str]]:
+        return self._vetted_calls
+
+    @vetted_calls.setter
+    def vetted_calls(self, functions: List[str]):
+        self._vetted_calls = functions
+
 
 class Stack(GenericSequence, Generic[T]):
     def __init__(self, initial_value: Iterable[T] = ()):