[TorchFix] Add codemod for unsafe load (#4715)

pytorch/test-infra#4671 added linter-only
`TorchUnsafeLoadVisitor`, but it turned out that the issue is so
widespread that manual fixes would be tedious.

The codemod is somewhat unsafe correctness-wise because full pickling
functionality may still be needed even without `pickle_module`, but I
think it's OK because it fixes a security-related issue and the codemods
need to be verified anyway.

Maybe later we should add something like Ruff's recently added
`--unsafe-fixes`: https://docs.astral.sh/ruff/linter/#fix-safety

I used this for pytorch/vision#8105

Author: kit1980

tests/fixtures/security/checker/load.py

@@ -1,7 +1,11 @@
 import random
+import dill
 import torch
+
 # Not OK
 torch.load('tensors.pt')
+torch.load('f.pt', pickle_module=dill, encoding='utf-8')
+
 # All these are OK
 torch.load('tensors.pt', weights_only=True)
 torch.load('tensors.pt', weights_only=False)

tests/fixtures/security/checker/load.txt

@@ -1 +1,2 @@
-4:1 TOR102 `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
+6:1 TOR102 `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.
+7:1 TOR102 `torch.load` without `weights_only` parameter is unsafe. Explicitly set `weights_only` to False only if you trust the data you load and full pickle functionality is needed, otherwise set `weights_only=True`.

torchfix/visitors/security/__init__.py

@@ -25,13 +25,31 @@ def visit_Call(self, node):
                     cst.metadata.WhitespaceInclusivePositionProvider, node
                 )
 
+                # Add `weights_only=True` if there is no `pickle_module`.
+                # (do not add `weights_only=False` with `pickle_module`, as it
+                # needs to be an explicit choice).
+                #
+                # This codemod is somewhat unsafe correctness-wise
+                # because full pickling functionality may still be needed
+                # even without `pickle_module`,
+                # so the changes need to be verified/tested.
+                replacement = None
+                pickle_module_arg = self.get_specific_arg(node, "pickle_module", 2)
+                if pickle_module_arg is None:
+                    weights_only_arg = cst.ensure_type(
+                        cst.parse_expression("f(weights_only=True)"), cst.Call
+                    ).args[0]
+                    replacement = node.with_changes(
+                        args=node.args + (weights_only_arg,)
+                    )
+
                 self.violations.append(
                     LintViolation(
                         error_code=self.ERROR_CODE,
                         message=self.MESSAGE,
                         line=position_metadata.start.line,
                         column=position_metadata.start.column,
                         node=node,
-                        replacement=None,
+                        replacement=replacement,
                     )
                 )