From 71bcd07be3bacb22b40ce3411bf908a3916e265b Mon Sep 17 00:00:00 2001
From: Saurabh Misra <misra.saurabh1@gmail.com>
Date: Tue, 9 Jul 2024 20:38:36 -0700
Subject: [PATCH] - Fix is_zero_diff checker by using ast.unparse rather than
 ast.dump - If only the imports are different then it is still 0 diff - add
 test cases - Codeflash does not support Python 3.13 yet. I saw pydantic crash
 while installing codeflash on 3.13 python

---
 codeflash/code_utils/code_replacer.py | 12 +++++++-
 codeflash/version.py                  |  4 +--
 pyproject.toml                        |  2 +-
 tests/test_code_replacement.py        | 42 ++++++++++++++++++++++++++-
 4 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/codeflash/code_utils/code_replacer.py b/codeflash/code_utils/code_replacer.py
index fcc9fbb1c..df9f0628b 100644
--- a/codeflash/code_utils/code_replacer.py
+++ b/codeflash/code_utils/code_replacer.py
@@ -253,8 +253,18 @@ def replace_function_definitions_in_module(
         contextual_functions,
         project_root_path,
     )
-    if ast.dump(ast.parse(new_code)) == ast.dump(ast.parse(source_code)):
+    if is_zero_diff(source_code, new_code):
         return False
     with open(module_abspath, "w", encoding="utf8") as file:
         file.write(new_code)
     return True
+
+
+def is_zero_diff(original_code: str, new_code: str) -> bool:
+    def normalize_for_diff(tree: ast.AST):
+        tree.body = [node for node in tree.body if not isinstance(node, (ast.Import, ast.ImportFrom))]
+        return tree
+
+    original_code_unparsed = ast.unparse(normalize_for_diff(ast.parse(original_code)))
+    new_code_unparsed = ast.unparse(normalize_for_diff(ast.parse(new_code)))
+    return original_code_unparsed == new_code_unparsed
diff --git a/codeflash/version.py b/codeflash/version.py
index 3d296a5e2..0275e0c0d 100644
--- a/codeflash/version.py
+++ b/codeflash/version.py
@@ -1,3 +1,3 @@
 # These version placeholders will be replaced by poetry-dynamic-versioning during `poetry build`.
-__version__ = "0.6.11"
-__version_tuple__ = (0, 6, 11)
+__version__ = "0.6.12"
+__version_tuple__ = (0, 6, 12)
diff --git a/pyproject.toml b/pyproject.toml
index 665caa38f..9eea44e9f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,7 @@ keywords = ["codeflash", "performance", "optimization", "ai", "code", "machine l
 # poetry self add poetry-dynamic-versioning
 
 [tool.poetry.dependencies]
-python = ">=3.9,<4.0"
+python = ">=3.9,<3.13"
 unidiff = ">=0.7.4"
 pytest = ">=7.0.0"
 gitpython = ">=3.1.31"
diff --git a/tests/test_code_replacement.py b/tests/test_code_replacement.py
index fe741f411..0b7ef2f3f 100644
--- a/tests/test_code_replacement.py
+++ b/tests/test_code_replacement.py
@@ -6,7 +6,11 @@ from argparse import Namespace
 from collections import defaultdict
 from pathlib import Path
 
-from codeflash.code_utils.code_replacer import replace_functions_and_add_imports, replace_functions_in_file
+from codeflash.code_utils.code_replacer import (
+    is_zero_diff,
+    replace_functions_and_add_imports,
+    replace_functions_in_file,
+)
 from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize
 from codeflash.optimization.optimizer import Optimizer
 
@@ -1499,3 +1503,39 @@ def cosine_similarity_top_k(
     return ret_idxs, scores
 '''
     )
+
+
+def test_0_diff_code_replacement():
+    original_code = """from __future__ import annotations
+    
+import numpy as np
+def functionA():
+    return np.array([1, 2, 3])
+"""
+    optim_code_a = """from __future__ import annotations
+import numpy as np
+def functionA():
+    return np.array([1, 2, 3])"""
+
+    assert is_zero_diff(original_code, optim_code_a)
+
+    optim_code_b = """
+import numpy as np
+def functionA():
+    return np.array([1, 2, 3])"""
+
+    assert is_zero_diff(original_code, optim_code_b)
+
+    optim_code_c = """
+def functionA():
+    return np.array([1, 2, 3])"""
+
+    assert is_zero_diff(original_code, optim_code_c)
+
+    optim_code_d = """from __future__ import annotations
+    
+import numpy as np
+def functionA():
+    return np.array([1, 2, 3, 4])
+"""
+    assert not is_zero_diff(original_code, optim_code_d)