Add pyproject and precommit configurations for black, isort and mypy and reformat the code accordingly

Author: UniverseFly

.pre-commit-config.yaml

@@ -0,0 +1,22 @@
+files: "^src/"
+
+repos:
+  - repo: https://github.com/psf/black
+    rev: 22.12.0
+    hooks:
+      - id: black
+        # It is recommended to specify the latest version of Python
+        # supported by your project here, or alternatively use
+        # pre-commit's default_language_version, see
+        # https://pre-commit.com/#top_level-default_language_version
+        language_version: python3.10
+  - repo: https://github.com/pycqa/isort
+    rev: 5.11.2
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+        name: isort (python)
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: "v0.991"  # Use the sha / tag you want to point at
+    hooks:
+      - id: mypy

calc_time.py

@@ -1,8 +1,9 @@
 import json
 import sys
 from pathlib import Path
-from matplotlib import pyplot as plt
+
 import numpy as np
+from matplotlib import pyplot as plt
 
 folder = Path(sys.argv[1])
 all_files = folder.glob("**/*.json")

evaluate.py

@@ -1,9 +1,10 @@
-from matplotlib import pyplot as plt
 import json
 import sys
-from pathlib import Path
-from typing import TypedDict, List
 from argparse import ArgumentParser
+from pathlib import Path
+from typing import List, TypedDict
+
+from matplotlib import pyplot as plt
 from tqdm import tqdm
 
 

pyproject.toml

@@ -0,0 +1,18 @@
+[tool.black]
+include = '\.pyi?$'
+line-length = 88
+target-version = ["py310"]
+
+[tool.isort]
+line_length = 88
+profile = "black"
+skip_gitignore = true
+
+[tool.mypy]
+check_untyped_defs = true
+follow_imports = "silent"
+ignore_missing_imports = true
+modules = ["main"]
+mypy_path = "src"
+packages = ["realm"]
+python_version = "3.10"

requirements.txt

@@ -1,19 +1,15 @@
 import itertools
-import pickle
 import subprocess
 from multiprocessing import Process
 from multiprocessing.connection import Connection
-import os
 from os import PathLike
 from pathlib import Path
 from typing import Any, Dict, List, Optional, cast
 
-import torch
-
+from realm import utils
 from realm.generation_defs import GenerationContext, Memorization
 from realm.lsp import LSPClient, TextFile, spec
 from realm.model import CodeT5ForRealm
-from realm import utils
 
 TIMEOUT_THRESHOULD = 300
 

src/realm/analyze/__init__.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
-from pathlib import Path
 from enum import Enum, auto
-from .generation_defs import LMInferenceConfig
-from .utils import JsonSerializable
+from pathlib import Path
 from typing import Any
 
+from .utils import JsonSerializable
+
 
 @dataclass(frozen=True)
 class MetaConfig(JsonSerializable):
@@ -74,6 +74,31 @@ def from_json(cls, d: str) -> "SynthesisMethod":
 SYNTHESIS_METHOD_REV_MAP = {value: key for key, value in SYNTHESIS_METHOD_MAP.items()}
 
 
+@dataclass(frozen=True)
+class LMInferenceConfig(JsonSerializable):
+    batch_size: int
+    temperature: float
+    top_k: int
+    max_new_tokens: int
+
+    def to_json(self) -> Any:
+        return {
+            "batch_size": self.batch_size,
+            "temperature": self.temperature,
+            "top_k": self.top_k,
+            "max_new_tokens": self.max_new_tokens,
+        }
+
+    @classmethod
+    def from_json(cls, d: dict) -> "LMInferenceConfig":
+        return LMInferenceConfig(
+            int(d["batch_size"]),
+            float(d["temperature"]),
+            int(d["top_k"]),
+            int(d["max_new_tokens"]),
+        )
+
+
 @dataclass(frozen=True)
 class RepairConfig(JsonSerializable):
     n_samples: int

src/realm/analyze/java_syntax.py

@@ -1,19 +1,19 @@
-from functools import partial
+import csv
 import itertools
+import multiprocessing as mp
+import os
+import subprocess
+from functools import partial
 from os import PathLike
 from pathlib import Path
-from typing import Dict, Iterator, List, NamedTuple
-from unidiff import PatchSet, PatchedFile
+from typing import Dict, Iterable, Iterator, List, NamedTuple, TypeVar
+
+import git
+from unidiff import PatchedFile, PatchSet
 from unidiff.patch import Line
+
 from realm import utils
-import subprocess
-import multiprocessing as mp
-import csv
-from pathlib import Path
-from typing import Dict, Iterable, Iterator, List, TypeVar
 from realm.utils import chunked
-import git
-import os
 
 Metadata = Dict[str, List[Dict[str, str]]]