New find-and-fix and sonar codemod for floating point equality (#450)

* new codemod to fix float equality

* add testing edge cases

* register and test float eq

* document codemod

* add sonar fix float equality codemod

* fix small bug

* Apply suggestions from code review

Co-authored-by: Dan D'Avella <[email protected]>

* Update src/core_codemods/fix_float_equality.py

Co-authored-by: Dan D'Avella <[email protected]>

---------

Co-authored-by: Dan D'Avella <[email protected]>

Author: clavedeluna

integration_tests/sonar/test_sonar_fix_float_equality.py

@@ -0,0 +1,29 @@
+from codemodder.codemods.test import SonarIntegrationTest
+from core_codemods.fix_float_equality import FixFloatEqualityTransformer
+from core_codemods.sonar.sonar_fix_float_equality import SonarFixFloatEquality
+
+
+class TestSonarFixFloatEquality(SonarIntegrationTest):
+    codemod = SonarFixFloatEquality
+    code_path = "tests/samples/fix_float_equality.py"
+    replacement_lines = [
+        (1, "import math\n"),
+        (2, "\n"),
+        (3, "def foo(a, b):\n"),
+        (4, "    return math.isclose(a, b - 0.1, rel_tol=1e-09, abs_tol=0.0)\n"),
+    ]
+    # fmt: off
+    expected_diff = (
+        """--- \n"""
+        """+++ \n"""
+        """@@ -1,2 +1,4 @@\n"""
+        """+import math\n"""
+        """+\n"""
+        """ def foo(a, b):\n"""
+        """-    return a == b - 0.1\n"""
+        """+    return math.isclose(a, b - 0.1, rel_tol=1e-09, abs_tol=0.0)\n"""
+    )
+    # fmt: on
+    expected_line_change = "2"
+    change_description = FixFloatEqualityTransformer.change_description
+    num_changes = 1

integration_tests/test_fix_float_equality.py

@@ -0,0 +1,33 @@
+from codemodder.codemods.test import BaseIntegrationTest
+from core_codemods.fix_float_equality import (
+    FixFloatEquality,
+    FixFloatEqualityTransformer,
+)
+
+
+class TestFixFloatEquality(BaseIntegrationTest):
+    codemod = FixFloatEquality
+    original_code = """
+    def foo(a, b):
+        return a == b - 0.1
+    """
+    expected_new_code = """
+    import math
+    
+    def foo(a, b):
+        return math.isclose(a, b - 0.1, rel_tol=1e-09, abs_tol=0.0)
+    """
+    # fmt: off
+    expected_diff = (
+        """--- \n"""
+        """+++ \n"""
+        """@@ -1,2 +1,4 @@\n"""
+        """+import math\n"""
+        """+\n"""
+        """ def foo(a, b):\n"""
+        """-    return a == b - 0.1\n"""
+        """+    return math.isclose(a, b - 0.1, rel_tol=1e-09, abs_tol=0.0)"""
+    )
+    # fmt: on
+    expected_line_change = "2"
+    change_description = FixFloatEqualityTransformer.change_description

src/codemodder/scripts/generate_docs.py

@@ -255,6 +255,10 @@ class DocMetadata:
         importance="Medium",
         guidance_explained="This change is safe and will prevent errors when calling on these instance or class methods..",
     ),
+    "fix-float-equality": DocMetadata(
+        importance="Medium",
+        guidance_explained="This change makes your code more accurate but in some cases it may be necessary to adjust the `abs_tol` and `rel_tol` parameter values for your particular calculations.",
+    ),
 }
 DEFECTDOJO_CODEMODS = {
     "django-secure-set-cookie": DocMetadata(
@@ -284,6 +288,7 @@ class DocMetadata:
     "secure-random-S2245",
     "enable-jinja2-autoescape-S5247",
     "url-sandbox-S5144",
+    "fix-float-equality-S1244",
 ]
 SONAR_CODEMODS = {
     name: DocMetadata(

src/core_codemods/__init__.py

@@ -20,6 +20,7 @@
 from .fix_deprecated_abstractproperty import FixDeprecatedAbstractproperty
 from .fix_deprecated_logging_warn import FixDeprecatedLoggingWarn
 from .fix_empty_sequence_comparison import FixEmptySequenceComparison
+from .fix_float_equality import FixFloatEquality
 from .fix_hasattr_call import TransformFixHasattrCall
 from .fix_missing_self_or_cls import FixMissingSelfOrCls
 from .fix_mutable_params import FixMutableParams
@@ -54,6 +55,7 @@
 from .sonar.sonar_enable_jinja2_autoescape import SonarEnableJinja2Autoescape
 from .sonar.sonar_exception_without_raise import SonarExceptionWithoutRaise
 from .sonar.sonar_fix_assert_tuple import SonarFixAssertTuple
+from .sonar.sonar_fix_float_equality import SonarFixFloatEquality
 from .sonar.sonar_fix_missing_self_or_cls import SonarFixMissingSelfOrCls
 from .sonar.sonar_flask_json_response_type import SonarFlaskJsonResponseType
 from .sonar.sonar_jwt_decode_verify import SonarJwtDecodeVerify
@@ -131,6 +133,7 @@
         FixEmptySequenceComparison,
         RemoveAssertionInPytestRaises,
         FixAssertTuple,
+        FixFloatEquality,
         LazyLogging,
         StrConcatInSeqLiteral,
         FixAsyncTaskInstantiation,
@@ -158,6 +161,7 @@
         SonarSecureRandom,
         SonarEnableJinja2Autoescape,
         SonarUrlSandbox,
+        SonarFixFloatEquality,
     ],
 )
 

src/core_codemods/docs/pixee_python_fix-float-equality.md

@@ -0,0 +1,12 @@
+In most programming languages, floating point arithmetic is imprecise due to the way floating point numbers are stored as binary representations. Moreover, the result of calculations with floats can vary based on when rounding happens. Using equality or inequality to compare floats or their operations will almost always be imprecise and lead to bugs.
+
+For these reasons, this codemod changes any operations involving equality or inequality with floats to the recommended `math.isclose` function. This codemod uses the default parameter values `rel_tol=1e-09` and  `abs_tol=0.0` but makes them explicit as a starting point for you to consider depending on your calculation needs.
+
+Our changes look like the following:
+```diff
++import math
++
+ def foo(a, b):
+-    return a == b - 0.1
++    return math.isclose(a, b - 0.1, rel_tol=1e-09, abs_tol=0.0)
+```

src/core_codemods/fix_float_equality.py

@@ -0,0 +1,98 @@
+import libcst as cst
+
+from codemodder.codemods.libcst_transformer import (
+    LibcstResultTransformer,
+    LibcstTransformerPipeline,
+)
+from codemodder.codemods.utils_mixin import NameAndAncestorResolutionMixin
+from core_codemods.api import Metadata, Reference, ReviewGuidance
+from core_codemods.api.core_codemod import CoreCodemod
+
+
+class FixFloatEqualityTransformer(
+    LibcstResultTransformer, NameAndAncestorResolutionMixin
+):
+    change_description = "Replace `==` or `!=` with `math.isclose`"
+
+    def leave_Comparison(
+        self, original_node: cst.Comparison, updated_node: cst.Comparison
+    ) -> cst.BaseExpression:
+        if not self.node_is_selected(original_node):
+            return updated_node
+
+        match original_node:
+            case cst.Comparison(
+                left=left, comparisons=[cst.ComparisonTarget() as target]
+            ):
+                if isinstance(
+                    target.operator, cst.Equal | cst.NotEqual
+                ) and self.at_least_one_float(left, right := target.comparator):
+                    self.add_needed_import("math")
+                    isclose_call = self.make_isclose_call(left, right)
+                    self.report_change(original_node)
+                    return (
+                        isclose_call
+                        if isinstance(target.operator, cst.Equal)
+                        else cst.UnaryOperation(
+                            operator=cst.Not(),
+                            expression=isclose_call,
+                        )
+                    )
+        return updated_node
+
+    def make_isclose_call(self, left, right):
+        return cst.Call(
+            func=cst.Attribute(
+                value=cst.Name(value="math"), attr=cst.Name(value="isclose")
+            ),
+            args=[
+                cst.Arg(value=left),
+                cst.Arg(value=right),
+                cst.Arg(
+                    keyword=cst.Name(value="rel_tol"),
+                    value=cst.Float(value="1e-09"),
+                    equal=cst.AssignEqual(
+                        whitespace_before=cst.SimpleWhitespace(""),
+                        whitespace_after=cst.SimpleWhitespace(""),
+                    ),
+                ),
+                cst.Arg(
+                    keyword=cst.Name(value="abs_tol"),
+                    value=cst.Float(value="0.0"),
+                    equal=cst.AssignEqual(
+                        whitespace_before=cst.SimpleWhitespace(""),
+                        whitespace_after=cst.SimpleWhitespace(""),
+                    ),
+                ),
+            ],
+        )
+
+    def at_least_one_float(self, left, right) -> bool:
+        left_type = self.resolve_expression(left)
+        right_type = self.resolve_expression(right)
+
+        match (left_type, right_type):
+            case (cst.Float(), _) | (_, cst.Float()):
+                return True
+            case (cst.BinaryOperation(), _):
+                return self.at_least_one_float(left_type.left, left_type.right)
+            case (_, cst.BinaryOperation()):
+                return self.at_least_one_float(right_type.left, right_type.right)
+        return False
+
+
+FixFloatEquality = CoreCodemod(
+    metadata=Metadata(
+        name="fix-float-equality",
+        summary="Use `math.isclose` Instead of Direct Equality for Floats",
+        review_guidance=ReviewGuidance.MERGE_AFTER_REVIEW,
+        references=[
+            Reference(
+                url="https://docs.python.org/3/tutorial/floatingpoint.html#floating-point-arithmetic-issues-and-limitations"
+            ),
+            Reference(url="https://docs.python.org/3/library/math.html#math.isclose"),
+        ],
+    ),
+    transformer=LibcstTransformerPipeline(FixFloatEqualityTransformer),
+    detector=None,
+)

src/core_codemods/sonar/sonar_fix_float_equality.py

@@ -0,0 +1,10 @@
+from core_codemods.fix_float_equality import FixFloatEquality
+from core_codemods.sonar.api import SonarCodemod
+
+SonarFixFloatEquality = SonarCodemod.from_core_codemod(
+    name="fix-float-equality-S1244",
+    other=FixFloatEquality,
+    rule_id="python:S1244",
+    rule_name="Floating point numbers should not be tested for equality",
+    rule_url="https://rules.sonarsource.com/python/type/Bug/RSPEC-1244/",
+)

tests/codemods/sonar/test_sonar_fix_float_equality.py

@@ -0,0 +1,45 @@
+import json
+
+from codemodder.codemods.test import BaseSASTCodemodTest
+from core_codemods.sonar.sonar_fix_float_equality import SonarFixFloatEquality
+
+
+class TestSonarFixFloatEquality(BaseSASTCodemodTest):
+    codemod = SonarFixFloatEquality
+    tool = "sonar"
+
+    def test_name(self):
+        assert self.codemod.name == "fix-float-equality-S1244"
+
+    def test_simple(self, tmpdir):
+        input_code = """
+        def foo(a, b):
+            return a == b - 0.1
+        """
+        expected_output = """
+        import math
+        
+        def foo(a, b):
+            return math.isclose(a, b - 0.1, rel_tol=1e-09, abs_tol=0.0)
+        """
+        issues = {
+            "issues": [
+                {
+                    "rule": "python:S1244",
+                    "status": "OPEN",
+                    "component": "code.py",
+                    "textRange": {
+                        "startLine": 3,
+                        "endLine": 3,
+                        "startOffset": 11,
+                        "endOffset": 23,
+                    },
+                }
+            ]
+        }
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected_output,
+            results=json.dumps(issues),
+        )