linting

Author: KRRT7

codeflash/code_utils/code_extractor.py

@@ -444,28 +444,28 @@ def resolve_star_import(module_name: str, project_root: Path) -> set[str]:
 
         if all_names is not None:
             return set(all_names)
-        else:
-            public_names = set()
-            for node in tree.body:
-                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
-                    if not node.name.startswith("_"):
-                        public_names.add(node.name)
-                elif isinstance(node, ast.Assign):
-                    for target in node.targets:
-                        if isinstance(target, ast.Name) and not target.id.startswith("_"):
-                            public_names.add(target.id)
-                elif isinstance(node, ast.AnnAssign):
-                    if isinstance(node.target, ast.Name) and not node.target.id.startswith("_"):
-                        public_names.add(node.target.id)
-                elif isinstance(node, ast.Import) or (
-                    isinstance(node, ast.ImportFrom) and not any(alias.name == "*" for alias in node.names)
-                ):
-                    for alias in node.names:
-                        name = alias.asname or alias.name
-                        if not name.startswith("_"):
-                            public_names.add(name)
 
-            return public_names
+        public_names = set()
+        for node in tree.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                if not node.name.startswith("_"):
+                    public_names.add(node.name)
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and not target.id.startswith("_"):
+                        public_names.add(target.id)
+            elif isinstance(node, ast.AnnAssign):
+                if isinstance(node.target, ast.Name) and not node.target.id.startswith("_"):
+                    public_names.add(node.target.id)
+            elif isinstance(node, ast.Import) or (
+                isinstance(node, ast.ImportFrom) and not any(alias.name == "*" for alias in node.names)
+            ):
+                for alias in node.names:
+                    name = alias.asname or alias.name
+                    if not name.startswith("_"):
+                        public_names.add(name)
+
+        return public_names  # noqa: TRY300
 
     except Exception as e:
         logger.warning(f"Error resolving star import for {module_name}: {e}")

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -13,15 +13,14 @@
 import dill as pickle
 
 
-class VerificationType(str, Enum): # moved from codeflash/verification/codeflash_capture.py
+class VerificationType(str, Enum):  # moved from codeflash/verification/codeflash_capture.py
     FUNCTION_CALL = (
         "function_call"  # Correctness verification for a test function, checks input values and output values)
     )
     INIT_STATE_FTO = "init_state_fto"  # Correctness verification for fto class instance attributes after init
     INIT_STATE_HELPER = "init_state_helper"  # Correctness verification for helper class instance attributes after init
 
 
-
 F = TypeVar("F", bound=Callable[..., Any])
 
 
@@ -131,9 +130,9 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
         exception = None
         gc.disable()
         try:
-            ret = func(*args, **kwargs) # coroutine creation has some overhead, though it is very small
+            ret = func(*args, **kwargs)  # coroutine creation has some overhead, though it is very small
             counter = time.perf_counter_ns()
-            return_value = await ret # let's measure the actual execution time of the code
+            return_value = await ret  # let's measure the actual execution time of the code
             codeflash_duration = time.perf_counter_ns() - counter
         except Exception as e:
             codeflash_duration = time.perf_counter_ns() - counter