fix(B614): suppress false positive on non-literal weights_only; add torch.jit.load regression coverage

bandit/plugins/pytorch_load.py

@@ -19,6 +19,13 @@
 With `weights_only=True`, PyTorch enforces a strict type check, ensuring
 that only torch.Tensor objects are loaded.
 
+Note: ``torch.jit.load`` uses TorchScript serialization (not pickle) and
+is **not** flagged by this check regardless of arguments.
+
+Note: When ``weights_only`` is supplied as a variable or non-literal
+expression, the check cannot statically resolve its value and will not
+raise a finding (benefit of the doubt, consistent with B615 behaviour).
+
 :Example:
 
 .. code-block:: none
@@ -42,6 +49,8 @@
 .. versionadded:: 1.7.10
 
 """
+import ast
+
 import bandit
 from bandit.core import issue
 from bandit.core import test_properties as test
@@ -56,13 +65,35 @@ def pytorch_load(context):
     `torch.serialization.load` with untrusted data can lead to
     arbitrary code execution. The safe alternative is to use
     `weights_only=True` or the safetensors library.
+
+    ``torch.jit.load`` is excluded because it uses TorchScript
+    serialization, not pickle, and therefore does not deserialize
+    arbitrary Python objects.
+
+    When ``weights_only`` is a non-literal expression (variable,
+    attribute, subscript, etc.) the value cannot be resolved at
+    static-analysis time; the check is suppressed to avoid false
+    positives, consistent with the treatment of similar arguments in
+    other B6xx checks.
     """
     imported = context.is_module_imported_exact("torch")
     qualname = context.call_function_name_qual
     if not imported and isinstance(qualname, str):
         return
 
     if qualname in {"torch.load", "torch.serialization.load"}:
+        # If weights_only is a non-literal expression (Name, Attribute,
+        # Subscript, Call, …) we cannot resolve its value statically.
+        # Give the caller the benefit of the doubt rather than raising a
+        # false positive.
+        call_node = context._context.get("call")
+        if call_node is not None:
+            for kw in getattr(call_node, "keywords", []):
+                if kw.arg == "weights_only" and not isinstance(
+                    kw.value, ast.Constant
+                ):
+                    return
+
         # For torch.load, check if weights_only=True is specified
         weights_only = context.get_call_arg_value("weights_only")
         if weights_only == "True" or weights_only is True:

examples/pytorch_load.py

@@ -28,3 +28,16 @@
 # Example of a torch.*.load call that should NOT trigger B614
 # Only pickle deserializers should trigger B614
 torch.utils.cpp_extension.load(name="example_ext", sources=[])
+
+# torch.jit.load uses TorchScript serialization, not pickle.
+# It has no weights_only parameter and should NOT trigger B614.
+jit_model = torch.jit.load('script_model.pt')
+
+# torch.jit.load with map_location is still TorchScript; should NOT
+# trigger B614.
+jit_model_cpu = torch.jit.load('script_model.pt', map_location='cpu')
+
+# weights_only passed as a variable (non-literal).
+# B614 cannot resolve the value statically; should NOT trigger B614.
+_weights_flag = True
+model_var = torch.load('model_weights.pth', weights_only=_weights_flag)