Add pre-commit hooks add lint docs (#51)

fix #48
* add pre-commit hooks
* add black to CI
* add contributor guidelines re: linting
* reformat code using black
* address mypy warnings re: generators

### Testing:

1.
```
pip install -r requirements-dev.txt
pre-commit install
pre-commit run --all-files
```
<img width="582" alt="image"
src="https://github.com/pytorch-labs/torchfix/assets/108101595/734960dd-4ca2-4d8a-b58e-dc917332b7f9">


2. 
```
git commit -m 'test'
```

3.
Black CI works:
https://github.com/pytorch-labs/torchfix/actions/runs/8791096314/job/24124561579?pr=51

.flake8

@@ -3,3 +3,4 @@ exclude = ./tests/fixtures/
 
 # Match black tool's default.
 max-line-length = 88
+extend-ignore = E203

.github/workflows/test-torchfix.yml

@@ -25,3 +25,6 @@ jobs:
       - name: Run mypy
         run: |
           mypy .
+      - name: Run black
+        run: |
+          black --check .

.pre-commit-config.yaml

@@ -0,0 +1,23 @@
+repos:
+  - repo: local
+    hooks:
+      - id: black
+        name: black
+        entry: black
+        language: system
+        types: [python]
+        args: ["--config=./pyproject.toml"]
+        exclude: ^tests/fixtures/
+      - id: flake8
+        name: flake8
+        entry: flake8
+        language: system
+        types: [python]
+        args: ["--config=./.flake8"]
+        exclude: ^tests/fixtures/
+      - id: mypy
+        name: mypy
+        entry: mypy
+        language: system
+        types: [python]
+        exclude: ^tests/fixtures/

CONTRIBUTING.md

@@ -14,10 +14,34 @@ We actively welcome your pull requests.
 1. Fork the repo and create your branch from `main`.
 2. If you've added code that should be tested, add tests.
 3. If you've changed APIs, update the documentation.
-4. Ensure the test suite passes.
-5. Make sure your code lints.
+4. Ensure the test suite passes (`pytest tests`).
+5. Make sure your code lints (see Linting section below).
 6. If you haven't already, complete the Contributor License Agreement ("CLA").
 
+## Linting
+
+We use `black`, `flake8`, and `mypy` to lint the code.
+```
+pip install -r requirements-dev.txt
+```
+
+Linting via pre-commit hook:
+```
+# install pre-commit hooks for the first time
+pre-commit install
+
+# manually run pre-commit hooks on all files (runs all linters)
+pre-commit run --all-files
+```
+
+Manually running individual linters:
+```
+black .
+flake8
+mypy .
+```
+
+
 ## Contributor License Agreement ("CLA")
 
 In order to accept your pull request, we need you to submit a CLA. You only

requirements-dev.txt

@@ -3,3 +3,5 @@ pytest==7.2.0
 libcst==1.1.0
 types-PyYAML==6.0.7
 mypy==1.7.0
+black==24.4.0
+pre-commit==3.7.0

torchfix/common.py

@@ -78,15 +78,17 @@ def get_specific_arg(
 
     @staticmethod
     def has_specific_arg(
-            node: cst.Call, arg_name: str, position: Optional[int] = None
+        node: cst.Call, arg_name: str, position: Optional[int] = None
     ) -> bool:
         """
         Check if the specific argument is present in a call.
         """
-        return TorchVisitor.get_specific_arg(
-            node, arg_name,
-            position if position is not None else -1
-        ) is not None
+        return (
+            TorchVisitor.get_specific_arg(
+                node, arg_name, position if position is not None else -1
+            )
+            is not None
+        )
 
     def add_violation(
         self,

torchfix/torchfix.py

@@ -48,7 +48,7 @@
 def GET_ALL_ERROR_CODES():
     codes = set()
     for cls in ALL_VISITOR_CLS:
-        codes |= set(error.error_code for error in cls.ERRORS)
+        codes |= {error.error_code for error in cls.ERRORS}
     return codes
 
 
@@ -83,7 +83,7 @@ def get_visitors_with_error_codes(error_codes):
         # only correspond to one visitor.
         found = False
         for visitor_cls in ALL_VISITOR_CLS:
-            if error_code in list(err.error_code for err in visitor_cls.ERRORS):
+            if any(error_code == err.error_code for err in visitor_cls.ERRORS):
                 visitor_classes.add(visitor_cls)
                 found = True
                 break

torchfix/visitors/misc/__init__.py

@@ -59,9 +59,11 @@ class TorchReentrantCheckpointVisitor(TorchVisitor):
     ]
 
     def visit_Call(self, node):
-        if (self.get_qualified_name_for_call(node) ==
-                "torch.utils.checkpoint.checkpoint" and
-                not self.has_specific_arg(node, "use_reentrant")):
+        if self.get_qualified_name_for_call(
+            node
+        ) == "torch.utils.checkpoint.checkpoint" and not self.has_specific_arg(
+            node, "use_reentrant"
+        ):
             # This codemod maybe  unsafe correctness-wise
             # if reentrant behavior is actually needed,
             # so the changes need to be verified/tested.

torchfix/visitors/nonpublic/__init__.py

@@ -19,13 +19,15 @@ class TorchNonPublicAliasVisitor(TorchVisitor):
 
     ERRORS: List[TorchError] = [
         TorchError(
-            "TOR104", (
+            "TOR104",
+            (
                 "Use of non-public function `{qualified_name}`, "
                 "please use `{public_name}` instead"
             ),
         ),
         TorchError(
-            "TOR105", (
+            "TOR105",
+            (
                 "Import of non-public function `{qualified_name}`, "
                 "please use `{public_name}` instead"
             ),

torchfix/visitors/security/__init__.py

@@ -15,15 +15,17 @@ class TorchUnsafeLoadVisitor(TorchVisitor):
             (
                 "`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,"
+                "the data you load "
+                "and full pickle functionality is needed,"
                 " otherwise set `weights_only=True`."
             ),
         )
     ]
 
     def visit_Call(self, node):
-        if self.get_qualified_name_for_call(node) == "torch.load" and \
-                not self.has_specific_arg(node, "weights_only"):
+        if self.get_qualified_name_for_call(
+            node
+        ) == "torch.load" and not self.has_specific_arg(node, "weights_only"):
             # 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).
@@ -37,9 +39,7 @@ def visit_Call(self, node):
                 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,)
-                )
+                replacement = node.with_changes(args=node.args + (weights_only_arg,))
             self.add_violation(
                 node,
                 error_code=self.ERRORS[0].error_code,