fix lint

Author: EYH0602

.pylintrc

@@ -36,3 +36,52 @@ dependencies = [
 [build-system]
 requires = ["uv_build>=0.8.10,<0.9.0"]
 build-backend = "uv_build"
+
+[tool.pylint.messages_control]
+disable = [
+    "raw-checker-failed",
+    "bad-inline-option",
+    "locally-disabled",
+    "file-ignored",
+    "suppressed-message",
+    "useless-suppression",
+    "deprecated-pragma",
+    "use-symbolic-message-instead",
+    "logging-fstring-interpolation",
+    "unused-import",
+    "ungrouped-imports",
+    "subprocess-run-check",
+    "too-many-instance-attributes",
+    "missing-module-docstring",
+    "unnecessary-ellipsis",          # pylint cannot work with typed Protocol,
+    "too-many-arguments",
+    "too-many-locals",
+    "too-many-positional-arguments",
+    "too-many-statements",
+    "unspecified-encoding",
+    "missing-class-docstring",
+]
+fail-under = 9
+max-line-length = 120
+
+[tool.pylint.typecheck]
+generated-members = ["cv2.*", "skimage.metrics.*"]
+
+[tool.mypy]
+python_version = "3.11"
+warn_return_any = true
+warn_unused_configs = true
+ignore_missing_imports = true
+check_untyped_defs = true
+allow_redefinition = false
+ignore_errors = false
+implicit_reexport = false
+local_partial_types = true
+no_implicit_optional = true
+strict_equality = true
+strict_optional = true
+warn_no_return = true
+warn_redundant_casts = true
+warn_unreachable = true
+warn_unused_ignores = false
+plugins = ["returns.contrib.mypy.returns_plugin"]

mypy.ini

@@ -1,17 +1,18 @@
+import json
+from typing import Iterable
+
 import fire
-import funcy
 from funcy_chain import Chain
-import logging
-import json
 from dacite import from_dict
+
 from tfbench.common import extract_function_name
 from tfbench.hs_parser import HASKELL_LANGUAGE
 from tfbench.hs_parser.ast_util import AST
 from tfbench.common import BenchmarkTask
-from typing import Iterable
 
 
 def build_dependency_dict(tasks: list[BenchmarkTask]) -> dict[str, str]:
+    """dependency mapping from function names to their signatures"""
     return {
         fn_name: t.signature
         for t in tasks
@@ -48,6 +49,7 @@ def get_func_calls(task: BenchmarkTask) -> set[str]:
 
 
 def _is_input(code: str, call: str) -> bool:
+    """check if a function call is an input to the task"""
     inputs: list[list[str]] = (
         Chain(code.splitlines())
         .filter(lambda l: "=" in l)
@@ -60,6 +62,8 @@ def _is_input(code: str, call: str) -> bool:
 
 
 def add_dependencies(dependency_dict: dict[str, str]):
+    """add dependencies to benchmark tasks"""
+
     def add_for_task(task: BenchmarkTask) -> BenchmarkTask:
         calls: Iterable[str] = get_func_calls(task)
         calls = filter(lambda c: not _is_input(task.code, c), calls)
@@ -73,6 +77,7 @@ def main(
     input_file: str = "data/source/base-4.20.0.0.jsonl",
     output_file: str = "out.jsonl",
 ):
+    """add dependency to a task json file"""
     with open(input_file, "r") as fp:
         tasks: list[BenchmarkTask] = (
             Chain(fp.readlines())

pyproject.toml

@@ -1,10 +1,10 @@
 from dataclasses import dataclass, field
 import re
 import copy
-from funcy import lmap
 import sys
 import io
 
+from funcy import lmap
 import markdown_to_json
 
 # Default hyper-parameters
@@ -40,10 +40,16 @@ class BenchmarkTask:
 
 
 def extract_function_name(task: BenchmarkTask) -> str | None:
+    """get function name from a task
+    For example, given a task with signature 'foo :: Int -> Int',
+    it will return 'foo'.
+    """
     return task.signature.split("::")[0].strip()
 
 
 def clean_tab_spaces(task: BenchmarkTask) -> BenchmarkTask:
+    """remove tab spaces from the code"""
+
     def clean(s: str) -> str:
         return re.sub(r"[ \t]+", " ", s)
 
@@ -56,6 +62,7 @@ def clean(s: str) -> str:
 
 
 def remove_comments(code: str) -> str:
+    """remove Haskell comments"""
     # multi-line
     # code = re.sub(r"\{\-[\s\S]*?\-\}", "", code)
     code = re.sub(r"\{\-.*?\-\}", "", code, flags=re.DOTALL)
@@ -65,6 +72,7 @@ def remove_comments(code: str) -> str:
 
 
 def get_sys_prompt(pure: bool) -> str:
+    """helper function to get system prompt for pure version of TF-Bench"""
     sys_prompt = SYSTEM_PROMPT + INSTRUCT_PROMPT
     sys_prompt += PURE_PROMPT if pure else CORE_PROMPT
     return sys_prompt
@@ -99,9 +107,9 @@ def remove_return_type(sig: str) -> str:
         # Keep '->' and discard what follows
         new_type = type_part[: last_arrow_index + 2]
         return func_part + " " + new_type.rstrip()
-    else:
-        # No top-level arrow found; return as-is
-        return sig
+
+    # No top-level arrow found; return as-is
+    return sig
 
 
 def get_prompt(task: BenchmarkTask) -> str:

src/tfbench/add_dependency.py

@@ -1,25 +1,25 @@
 """main script to build Haskell dataset"""
 
+from enum import IntEnum
+import os
+import logging
+import json
+
 import fire
 from returns.result import Success, Failure
 from returns.io import IOResult, IOSuccess, IOFailure
 from funcy_chain import Chain
-import os
-import logging
-from enum import IntEnum
-import json
 from tqdm import tqdm
 from funcy import lmap
+
 from tfbench.hs_parser import HASKELL_LANGUAGE
 from tfbench.hs_parser.ast_util import AST, HaskellFunction
-from tfbench.hs_parser.polymorphism import get_polymorphic_type, PolymorphicType
+from tfbench.hs_parser.polymorphism import get_polymorphic_type
 from tfbench.common import remove_comments
 
 
 def wrap_repo(s: str) -> str:
-    # NOTE: this is a placeholder function
-    # the implementation depends on how the repo is downloaded
-    # please refer to HMTypes4Py repo
+    """NOTE: this is a placeholder function"""
     return s
 
 
@@ -44,6 +44,7 @@ def collect_hs_files(root: str):
 
 
 def collect_from_file(file_path: str) -> list[dict[str, str]]:
+    """extract benchmark tasks from a Haskell file"""
     with open(file_path, "r", errors="replace") as fp:
         code = fp.read()
 
@@ -66,6 +67,7 @@ def _to_json(func: HaskellFunction) -> dict[str, str]:
 def collect_from_repo(
     repo_id: str, repo_root: str, source_root: str
 ) -> IOResult[int, CollectionErrorCode]:
+    """Collect all tasks from a Haskell repository"""
     repo_path = os.path.join(repo_root, wrap_repo(repo_id))
     if not os.path.exists(repo_path) or not os.path.isdir(repo_path):
         return IOFailure(CollectionErrorCode.REPO_NOT_FOUND)
@@ -98,6 +100,7 @@ def main(
     repo_root: str = "data/repos",
     oroot: str = "data/source",
 ):
+    """script to extract benchmark tasks from Haskell repositories"""
     with open(input_repo_list_path) as fp:
         repo_id_list = [l.strip() for l in fp.readlines()]
 

src/tfbench/common.py

@@ -1,13 +1,15 @@
-import fire
 import json
 import logging
-from tfbench.common import BenchmarkTask
-from tfbench.postprocessing import postprocess, TASK_STRATEGIES, RESPONSE_STRATEGIES
-from funcy_chain import Chain
-from dacite import from_dict
 from itertools import starmap
 import re
 
+import fire
+from funcy_chain import Chain
+from dacite import from_dict
+
+from tfbench.common import BenchmarkTask
+from tfbench.postprocessing import postprocess, TASK_STRATEGIES, RESPONSE_STRATEGIES
+
 
 def tokenize_type_signature(sig: str) -> list[str]:
     """
@@ -64,6 +66,7 @@ def alpha_equiv(s1: str, s2: str) -> bool:
 
 
 def evaluate_one_task(task: BenchmarkTask, result: str) -> bool:
+    """evaluate a single task against its result by alpha equivalence"""
     ground_truth = postprocess(task.signature, TASK_STRATEGIES)
     result = postprocess(result, RESPONSE_STRATEGIES)
     return alpha_equiv(ground_truth, result)
@@ -72,6 +75,7 @@ def evaluate_one_task(task: BenchmarkTask, result: str) -> bool:
 def evaluate(
     benchmark_f: list[BenchmarkTask], results: list[str]
 ) -> dict[str, int | float]:
+    """evaluate all generation results"""
 
     assert len(benchmark_f) == len(results)
     eval_results = starmap(evaluate_one_task, zip(benchmark_f, results))
@@ -89,6 +93,7 @@ def main(
     benchmark_file: str = "Benchmark-F.jsonl",
     results_file: str = "data/experiment/gpt_generated_responses.jsonl",
 ):
+    """script to run all evaluation tasks"""
     with open(benchmark_file, "r") as file:
         benchmark_f: list[BenchmarkTask] = (
             Chain(file.readlines())

src/tfbench/dataset.py

@@ -2,10 +2,10 @@
 Experiment script for OpenAI models
 """
 
-from openai import OpenAI
-from anthropic import Anthropic, InternalServerError
 from typing import Callable
 
+from openai import OpenAI
+from anthropic import Anthropic, InternalServerError
 from google import genai
 from google.genai.types import GenerateContentConfig, ThinkingConfig
 
@@ -64,6 +64,8 @@ def get_oai_ttc_model(
     model: str = "o1-preview-2024-09-12",
     pure: bool = False,
 ) -> Callable[[str], str | None]:
+    """OpenAI Reasoning Models"""
+
     def generate_type_signature(prompt: str) -> str | None:
         completion = client.chat.completions.create(
             messages=[
@@ -83,6 +85,8 @@ def get_oai_model(
     model: str = "gpt-3.5-turbo",
     pure: bool = False,
 ) -> Callable[[str], str | None]:
+    """Regular OpenAI Models"""
+
     def generate_type_signature(prompt: str) -> str | None:
         completion = client.chat.completions.create(
             messages=[
@@ -107,6 +111,8 @@ def get_ant_ttc_model(
     pure: bool = False,
     thinking_budget: int = 1024,
 ) -> Callable[[str], str | None]:
+    "Claude Reasoning Models"
+
     def generate_type_signature(prompt: str) -> str | None:
         try:
             message = client.beta.messages.create(
@@ -138,6 +144,8 @@ def get_ant_model(
     model: str = "claude-3-5-sonnet-20240620",
     pure: bool = False,
 ) -> Callable[[str], str | None]:
+    """Claude regular Models"""
+
     def generate_type_signature(prompt: str) -> str | None:
         try:
             message = client.messages.create(
@@ -155,8 +163,8 @@ def generate_type_signature(prompt: str) -> str | None:
         if len(contents) > 0:
             text = contents[0].text  # type: ignore
             return text if isinstance(text, str) else None
-        else:
-            return None
+
+        return None
 
     return generate_type_signature
 
@@ -166,6 +174,8 @@ def get_gemini_model(
     model: str = "gemini-2.0-flash-lite",
     pure: bool = False,
 ) -> Callable[[str], str | None]:
+    """Gemini Models"""
+
     def generate_type_signature(prompt: str) -> str | None:
         response = client.models.generate_content(
             model=model,
@@ -185,6 +195,8 @@ def get_gemini_ttc_model(
     pure: bool = False,
     thinking_budget: int = 1024,
 ) -> Callable[[str], str | None]:
+    """Gemini Reasoning Models"""
+
     def generate_type_signature(prompt: str) -> str | None:
         response = client.models.generate_content(
             model=model,

src/tfbench/evaluation.py

@@ -2,8 +2,8 @@
 Experiment script for OSS models using Ollama
 """
 
-from ollama import Client as OllamaClient, ResponseError
 from typing import Union
+from ollama import Client as OllamaClient, ResponseError
 from tfbench.common import get_sys_prompt
 
 OLLAMA_OSS = [
@@ -61,15 +61,17 @@ def get_ollama_model(
     Parameters:
         client (OllamaClient): The Ollama client instance used for sending requests to the model.
 
-        model (str): Name of the model to use for generating type signatures. Must be one of the predefined models in OLLAMA_MODELS.
-                     Default is "llama3:8b".
+        model (str): Name of the model to use for generating type signatures.
+                    Must be one of the predefined models in OLLAMA_MODELS.
+                    Default is "llama3:8b".
 
         pure (bool): If True, uses the original variable naming in type inference.
                      If False, uses rewritten variable naming (e.g., `v1`, `v2`, ...). Default is False.
 
     Returns:
-        Callable[[str], Union[str, None]]: A function that takes a prompt string as input and returns the generated type
-                                           signature as a string, or None if the generation fails.
+        Callable[[str], Union[str, None]]:
+            A function that takes a prompt string as input and returns the generated type
+            signature as a string, or None if the generation fails.
     """
 
     def generate_type_signature(prompt: str) -> Union[str, None]: