minor update

Author: KRRT7

codeflash/verification/verification_utils.py

@@ -31,33 +31,34 @@ def delete_multiple_if_name_main(test_ast: ast.Module) -> ast.Module:
 
 
 class ModifyInspiredTests(ast.NodeTransformer):
-    """This isn't being used right now"""
+    """Transformer for modifying inspired test classes.
 
-    def __init__(self, import_list, test_framework):
+    Class is currently not in active use.
+    """
+
+    def __init__(self, import_list: list[ast.AST], test_framework: str) -> None:
         self.import_list = import_list
         self.test_framework = test_framework
 
-    def visit_Import(self, node: ast.Import):
+    def visit_Import(self, node: ast.Import) -> None:
         self.import_list.append(node)
 
-    def visit_ImportFrom(self, node: ast.ImportFrom):
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
         self.import_list.append(node)
 
-    def visit_ClassDef(self, node: ast.ClassDef):
+    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
         if self.test_framework != "unittest":
             return node
         found = False
         if node.bases:
             for base in node.bases:
-                if isinstance(base, ast.Attribute):
-                    if base.attr == "TestCase" and base.value.id == "unittest":
-                        found = True
-                        break
-                if isinstance(base, ast.Name):
-                    # TODO: Possibility that this is not a unittest.TestCase
-                    if base.id == "TestCase":
-                        found = True
-                        break
+                if isinstance(base, ast.Attribute) and base.attr == "TestCase" and base.value.id == "unittest":
+                    found = True
+                    break
+                # TODO: Check if this is actually a unittest.TestCase
+                if isinstance(base, ast.Name) and base.id == "TestCase":
+                    found = True
+                    break
         if not found:
             return node
         node.name = node.name + "Inspired"

codeflash/verification/verifier.py

@@ -71,7 +71,7 @@ def merge_unit_tests(unit_test_source: str, inspired_unit_tests: str, test_frame
     except SyntaxError as e:
         logger.exception(f"Syntax error in code: {e}")
         return unit_test_source
-    import_list: list[ast.stmt] = list()
+    import_list: list[ast.stmt] = []
     modified_ast = ModifyInspiredTests(import_list, test_framework).visit(inspired_unit_tests_ast)
     if test_framework == "pytest":
         # Because we only want to modify the top level test functions

pyproject.toml

@@ -176,13 +176,14 @@ ignore = [
     "TD003",
     "TD004",
     "PLR2004",
-    "UP007"
+    "UP007",
+    "N802", # we use a lot of stdlib which follows this convention
 ]
 
 [tool.ruff.lint.flake8-type-checking]
 strict = true
 runtime-evaluated-base-classes = ["pydantic.BaseModel"]
-runtime-evaluated-decorators = ["pydantic.validate_call"]
+runtime-evaluated-decorators = ["pydantic.validate_call", "pydantic.dataclasses.dataclass"]
 
 [tool.ruff.lint.pep8-naming]
 classmethod-decorators = [