Merge branch 'main' into codeflash-trace-decorator

Author: alvin-r

codeflash/verification/comparator.py

@@ -155,6 +155,8 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool:
 
         if HAS_PANDAS and isinstance(orig, (pandas.CategoricalDtype, pandas.Interval, pandas.Period)):
             return orig == new
+        if HAS_PANDAS and pandas.isna(orig) and pandas.isna(new):
+            return True
 
         # This should be at the end of all numpy checking
         try:
@@ -242,3 +244,7 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool:
         logger.error(f"RecursionError while comparing objects: {e}")
         sentry_sdk.capture_exception(e)
         return False
+    except Exception as e:
+        logger.error(f"Error while comparing objects: {e}")
+        sentry_sdk.capture_exception(e)
+        return False

tests/test_comparator.py

@@ -468,6 +468,36 @@ def test_pandas():
     assert comparator(an, ao)
     assert not comparator(an, ap)
 
+    assert comparator(pd.NA, pd.NA)
+    assert not comparator(pd.NA, None)
+    assert not comparator(None, pd.NA)
+
+    s1 = pd.Series([1, 2, pd.NA, 4])
+    s2 = pd.Series([1, 2, pd.NA, 4])
+    s3 = pd.Series([1, 2, None, 4])
+
+    assert comparator(s1, s2)
+    assert not comparator(s1, s3)
+
+    df1 = pd.DataFrame({'a': [1, 2, pd.NA], 'b': [4, pd.NA, 6]})
+    df2 = pd.DataFrame({'a': [1, 2, pd.NA], 'b': [4, pd.NA, 6]})
+    df3 = pd.DataFrame({'a': [1, 2, None], 'b': [4, None, 6]})
+    assert comparator(df1, df2)
+    assert not comparator(df1, df3)
+
+    d1 = {'a': pd.NA, 'b': [1, pd.NA, 3]}
+    d2 = {'a': pd.NA, 'b': [1, pd.NA, 3]}
+    d3 = {'a': None, 'b': [1, None, 3]}
+    assert comparator(d1, d2)
+    assert not comparator(d1, d3)
+
+    s1 = pd.Series([1, 2, pd.NA, 4])
+    s2 = pd.Series([1, 2, pd.NA, 4])
+
+    filtered1 = s1[s1 > 1]
+    filtered2 = s2[s2 > 1]
+    assert comparator(filtered1, filtered2)
+
 
 def test_pyrsistent():
     try: