perf

perf: optimize comparator with module-level constants and identity check
- Hoist _equality_types tuple and conditional library imports to module level
- Use isinstance for dict view type checks instead of string comparison
- Use isinstance for user-defined __eq__ detection instead of str(type(...))
- Add identity short-circuit (orig is new) at the top of comparator()

style: suppress false-positive PD011 on tf.SparseTensor.values

Author: KRRT7

codeflash/verification/comparator.py

@@ -28,6 +28,33 @@
 HAS_NUMBA = find_spec("numba") is not None
 HAS_PYARROW = find_spec("pyarrow") is not None
 
+if HAS_NUMPY:
+    import numpy as np
+if HAS_SCIPY:
+    import scipy  # type: ignore  # noqa: PGH003
+if HAS_JAX:
+    import jax  # type: ignore  # noqa: PGH003
+    import jax.numpy as jnp  # type: ignore  # noqa: PGH003
+if HAS_XARRAY:
+    import xarray  # type: ignore  # noqa: PGH003
+if HAS_TENSORFLOW:
+    import tensorflow as tf  # type: ignore  # noqa: PGH003
+if HAS_SQLALCHEMY:
+    import sqlalchemy  # type: ignore  # noqa: PGH003
+if HAS_PYARROW:
+    import pyarrow as pa  # type: ignore  # noqa: PGH003
+if HAS_PANDAS:
+    import pandas  # noqa: ICN001
+if HAS_TORCH:
+    import torch  # type: ignore  # noqa: PGH003
+if HAS_NUMBA:
+    import numba  # type: ignore  # noqa: PGH003
+    from numba.core.dispatcher import Dispatcher  # type: ignore  # noqa: PGH003
+    from numba.typed import Dict as NumbaDict  # type: ignore  # noqa: PGH003
+    from numba.typed import List as NumbaList  # type: ignore  # noqa: PGH003
+if HAS_PYRSISTENT:
+    import pyrsistent  # type: ignore  # noqa: PGH003
+
 # Pattern to match pytest temp directories: /tmp/pytest-of-<user>/pytest-<N>/
 # These paths vary between test runs but are logically equivalent
 PYTEST_TEMP_PATH_PATTERN = re.compile(r"/tmp/pytest-of-[^/]+/pytest-\d+/")  # noqa: S108
@@ -36,6 +63,31 @@
 # Created by tempfile.mkdtemp() or tempfile.TemporaryDirectory()
 PYTHON_TEMPFILE_PATTERN = re.compile(r"/tmp/tmp[a-zA-Z0-9_]+/")  # noqa: S108
 
+_DICT_KEYS_TYPE = type({}.keys())
+_DICT_VALUES_TYPE = type({}.values())
+_DICT_ITEMS_TYPE = type({}.items())
+
+_EQUALITY_TYPES = (
+    int,
+    bool,
+    complex,
+    type(None),
+    type(Ellipsis),
+    decimal.Decimal,
+    set,
+    bytes,
+    bytearray,
+    memoryview,
+    frozenset,
+    enum.Enum,
+    type,
+    range,
+    slice,
+    OrderedDict,
+    types.GenericAlias,
+    *((_union_type,) if (_union_type := getattr(types, "UnionType", None)) else ()),
+)
+
 
 def _normalize_temp_path(path: str) -> str:
     """Normalize temporary file paths by replacing session-specific components.
@@ -145,27 +197,7 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return _normalize_temp_path(orig) == _normalize_temp_path(new)
             return False
 
-        _equality_types = (
-            int,
-            bool,
-            complex,
-            type(None),
-            type(Ellipsis),
-            decimal.Decimal,
-            set,
-            bytes,
-            bytearray,
-            memoryview,
-            frozenset,
-            enum.Enum,
-            type,
-            range,
-            slice,
-            OrderedDict,
-            types.GenericAlias,
-            *((_union_type,) if (_union_type := getattr(types, "UnionType", None)) else ()),
-        )
-        if isinstance(orig, _equality_types):
+        if isinstance(orig, _EQUALITY_TYPES):
             return orig == new
         if isinstance(orig, float):
             if math.isnan(orig) and math.isnan(new):
@@ -184,9 +216,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
             return comparator(orig_referent, new_referent, superset_obj)
 
         if HAS_JAX:
-            import jax  # type: ignore  # noqa: PGH003
-            import jax.numpy as jnp  # type: ignore  # noqa: PGH003
-
             # Handle JAX arrays first to avoid boolean context errors in other conditions
             if isinstance(orig, jax.Array):
                 if orig.dtype != new.dtype:
@@ -197,15 +226,11 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
 
         # Handle xarray objects before numpy to avoid boolean context errors
         if HAS_XARRAY:
-            import xarray  # type: ignore  # noqa: PGH003
-
             if isinstance(orig, (xarray.Dataset, xarray.DataArray)):
                 return orig.identical(new)
 
         # Handle TensorFlow objects early to avoid boolean context errors
         if HAS_TENSORFLOW:
-            import tensorflow as tf  # type: ignore  # noqa: PGH003
-
             if isinstance(orig, tf.Tensor):
                 if orig.dtype != new.dtype:
                     return False
@@ -231,7 +256,9 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 if not comparator(orig.dense_shape.numpy(), new.dense_shape.numpy(), superset_obj):
                     return False
                 return comparator(orig.indices.numpy(), new.indices.numpy(), superset_obj) and comparator(
-                    orig.values.numpy(), new.values.numpy(), superset_obj
+                    orig.values.numpy(),  # noqa: PD011
+                    new.values.numpy(),  # noqa: PD011
+                    superset_obj,
                 )
 
             if isinstance(orig, tf.RaggedTensor):
@@ -242,8 +269,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return comparator(orig.to_list(), new.to_list(), superset_obj)
 
         if HAS_SQLALCHEMY:
-            import sqlalchemy  # type: ignore  # noqa: PGH003
-
             try:
                 insp = sqlalchemy.inspection.inspect(orig)
                 insp = sqlalchemy.inspection.inspect(new)
@@ -259,8 +284,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
             except sqlalchemy.exc.NoInspectionAvailable:
                 pass
 
-        if HAS_SCIPY:
-            import scipy  # type: ignore  # noqa: PGH003
         # scipy condition because dok_matrix type is also a instance of dict, but dict comparison doesn't work for it
         if isinstance(orig, dict) and not (HAS_SCIPY and isinstance(orig, scipy.sparse.spmatrix)):
             if superset_obj:
@@ -279,21 +302,14 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
             return comparator(dict(orig), dict(new), superset_obj)
 
         # Handle dict view types (dict_keys, dict_values, dict_items)
-        # Use type name checking since these are not directly importable types
-        type_name = type(orig).__name__
-        if type_name == "dict_keys":
-            # dict_keys can be compared as sets (order doesn't matter)
+        if isinstance(orig, _DICT_KEYS_TYPE):
             return comparator(set(orig), set(new))
-        if type_name == "dict_values":
-            # dict_values need element-wise comparison (order matters)
+        if isinstance(orig, _DICT_VALUES_TYPE):
             return comparator(list(orig), list(new))
-        if type_name == "dict_items":
-            # Convert to dict for order-insensitive comparison (handles unhashable values)
+        if isinstance(orig, _DICT_ITEMS_TYPE):
             return comparator(dict(orig), dict(new), superset_obj)
 
         if HAS_NUMPY:
-            import numpy as np
-
             if isinstance(orig, (np.datetime64, np.timedelta64)):
                 # Handle NaT (Not a Time) - numpy's equivalent of NaN for datetime
                 if np.isnat(orig) and np.isnat(new):
@@ -355,8 +371,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
             return (orig != new).nnz == 0
 
         if HAS_PYARROW:
-            import pyarrow as pa  # type: ignore  # noqa: PGH003
-
             if isinstance(orig, pa.Table):
                 if orig.schema != new.schema:
                     return False
@@ -399,8 +413,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return bool(orig.equals(new))
 
         if HAS_PANDAS:
-            import pandas  # noqa: ICN001
-
             if isinstance(
                 orig, (pandas.DataFrame, pandas.Series, pandas.Index, pandas.Categorical, pandas.arrays.SparseArray)
             ):
@@ -431,8 +443,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
             pass
 
         if HAS_TORCH:
-            import torch  # type: ignore  # noqa: PGH003
-
             if isinstance(orig, torch.Tensor):
                 if orig.dtype != new.dtype:
                     return False
@@ -451,11 +461,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return orig == new
 
         if HAS_NUMBA:
-            import numba
-            from numba.core.dispatcher import Dispatcher
-            from numba.typed import Dict as NumbaDict
-            from numba.typed import List as NumbaList
-
             # Handle numba typed List
             if isinstance(orig, NumbaList):
                 if len(orig) != len(new):
@@ -487,8 +492,6 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
                 return orig.py_func is new.py_func
 
         if HAS_PYRSISTENT:
-            import pyrsistent  # type: ignore  # noqa: PGH003
-
             if isinstance(
                 orig,
                 (
@@ -534,7 +537,7 @@ def comparator(orig: Any, new: Any, superset_obj: bool = False) -> bool:
         # If the object passed has a user defined __eq__ method, use that
         # This could fail if the user defined __eq__ is defined with C-extensions
         try:
-            if hasattr(orig, "__eq__") and str(type(orig.__eq__)) == "<class 'method'>":
+            if hasattr(orig, "__eq__") and isinstance(orig.__eq__, types.MethodType):
                 return orig == new
         except Exception:
             pass