Codemod to fix math.isclose comparison to 0 (#475)

* codemod for math.isclose without abs_tol

* add more test cases

* document codemod

* add sonar codemod for fix math.isclose

* document sonar

Author: clavedeluna

integration_tests/sonar/test_sonar_fix_math_isclose.py

@@ -0,0 +1,14 @@
+from codemodder.codemods.test import SonarIntegrationTest
+from core_codemods.sonar.sonar_fix_math_isclose import (
+    FixMathIsCloseSonarTransformer,
+    SonarFixMathIsClose,
+)
+
+
+class TestSonarFixMathIsClose(SonarIntegrationTest):
+    codemod = SonarFixMathIsClose
+    code_path = "tests/samples/fix_math_isclose.py"
+    replacement_lines = [(5, "    return math.isclose(a, 0, abs_tol=1e-09)\n")]
+    expected_diff = "--- \n+++ \n@@ -2,4 +2,4 @@\n \n \n def foo(a):\n-    return math.isclose(a, 0)\n+    return math.isclose(a, 0, abs_tol=1e-09)\n"
+    expected_line_change = "5"
+    change_description = FixMathIsCloseSonarTransformer.change_description

integration_tests/test_fix_math_isclose.py

@@ -0,0 +1,32 @@
+from codemodder.codemods.test import BaseIntegrationTest
+from core_codemods.fix_math_isclose import FixMathIsClose, FixMathIsCloseTransformer
+
+
+class TestFixMathIsClose(BaseIntegrationTest):
+    codemod = FixMathIsClose
+    original_code = """
+    import math
+
+    def foo(a):
+        return math.isclose(a, 0)
+    """
+    expected_new_code = """
+    import math
+
+    def foo(a):
+        return math.isclose(a, 0, abs_tol=1e-09)
+    """
+    # fmt: off
+    expected_diff = (
+        """--- \n"""
+        """+++ \n"""
+        """@@ -1,4 +1,4 @@\n"""
+        """ import math\n"""
+        """ \n"""
+        """ def foo(a):\n"""
+        """-    return math.isclose(a, 0)\n"""
+        """+    return math.isclose(a, 0, abs_tol=1e-09)"""
+    )
+    # fmt: on
+    expected_line_change = "4"
+    change_description = FixMathIsCloseTransformer.change_description

src/codemodder/codemods/utils.py

@@ -192,3 +192,27 @@ def is_assigned_to_True(original_node: cst.Assign):
         isinstance(original_node.value, cst.Name)
         and original_node.value.value == "True"
     )
+
+
+def is_zero(node: cst.CSTNode) -> bool:
+    match node:
+        case cst.Integer() | cst.Float():
+            try:
+                return float(node.value) == 0
+            except (ValueError, TypeError):
+                return False
+        case cst.Call(func=cst.Name("int"), args=[]) | cst.Call(
+            func=cst.Name("float"), args=[]
+        ):
+            # int() or float() == 0
+            return True
+        case cst.Call(
+            func=cst.Name("int"), args=[cst.Arg(value=cst.Name(value="False"))]
+        ) | cst.Call(
+            func=cst.Name("float"), args=[cst.Arg(value=cst.Name(value="False"))]
+        ):
+            # int(False) or float(False)
+            return True
+        case cst.Call(func=cst.Name("int")) | cst.Call(func=cst.Name("float")):
+            return is_zero(node.args[0].value)
+    return False

src/codemodder/scripts/generate_docs.py

@@ -259,6 +259,10 @@ class 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.",
     ),
+    "fix-math-isclose": 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` parameter value for your particular calculations.",
+    ),
 }
 DEFECTDOJO_CODEMODS = {
     "django-secure-set-cookie": DocMetadata(
@@ -289,6 +293,7 @@ class DocMetadata:
     "enable-jinja2-autoescape-S5247",
     "url-sandbox-S5144",
     "fix-float-equality-S1244",
+    "fix-math-isclose-S6727",
 ]
 SONAR_CODEMODS = {
     name: DocMetadata(

src/core_codemods/__init__.py

@@ -22,6 +22,7 @@
 from .fix_empty_sequence_comparison import FixEmptySequenceComparison
 from .fix_float_equality import FixFloatEquality
 from .fix_hasattr_call import TransformFixHasattrCall
+from .fix_math_isclose import FixMathIsClose
 from .fix_missing_self_or_cls import FixMissingSelfOrCls
 from .fix_mutable_params import FixMutableParams
 from .flask_enable_csrf_protection import FlaskEnableCSRFProtection
@@ -56,6 +57,7 @@
 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_math_isclose import SonarFixMathIsClose
 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
@@ -141,6 +143,7 @@
         TransformFixHasattrCall,
         FixDataclassDefaults,
         FixMissingSelfOrCls,
+        FixMathIsClose,
     ],
 )
 
@@ -162,6 +165,7 @@
         SonarEnableJinja2Autoescape,
         SonarUrlSandbox,
         SonarFixFloatEquality,
+        SonarFixMathIsClose,
     ],
 )
 

src/core_codemods/docs/pixee_python_fix-math-isclose.md

@@ -0,0 +1,12 @@
+The default value for the `abs_tol` argument to a `math.isclose` call is `0`. Using this default when comparing a value against `0`, such as in `math.isclose(a, 0)` is equivalent to a strict equality check to `0`, which is not the intended use of the `math.isclose` function.
+
+This codemod adds `abs_tol=1e-09` to any call to `math.isclose` with one of of the first arguments evaluating to `0` if `abs_tol` is not already specified. `1e-09` is a starting point for you to consider depending on your calculation needs.
+
+Our changes look like the following:
+```diff
++import math
++
+ def foo(a):
+-    return math.isclose(a, 0)
++    return math.isclose(a, 0, abs_tol=1e-09)
+```

src/core_codemods/fix_math_isclose.py

@@ -0,0 +1,63 @@
+import libcst as cst
+
+from codemodder.codemods.libcst_transformer import (
+    LibcstResultTransformer,
+    LibcstTransformerPipeline,
+    NewArg,
+)
+from codemodder.codemods.utils import is_zero
+from codemodder.codemods.utils_mixin import NameAndAncestorResolutionMixin
+from core_codemods.api import Metadata, Reference, ReviewGuidance
+from core_codemods.api.core_codemod import CoreCodemod
+
+
+class FixMathIsCloseTransformer(
+    LibcstResultTransformer,
+    NameAndAncestorResolutionMixin,
+):
+    change_description = "Add `abs_tol` to `math.isclose` call"
+
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
+        if (
+            not self.node_is_selected(original_node)
+            or self.find_base_name(original_node.func) != "math.isclose"
+            or len(original_node.args) < 2
+        ):
+            return updated_node
+
+        if self.at_least_one_zero_arg(original_node.args):
+            for arg in original_node.args[2:]:
+                match arg:
+                    case cst.Arg(keyword=cst.Name(value="abs_tol")) as matched_arg:
+                        # A `abs_tol` kwarg set to not 0 is acceptable if comparing to 0
+                        if not is_zero(matched_arg.value):
+                            return updated_node
+
+            new_args = self.replace_args(
+                original_node,
+                [NewArg(name="abs_tol", value="1e-09", add_if_missing=True)],
+            )
+
+            self.report_change(original_node)
+            return self.update_arg_target(updated_node, new_args)
+
+        return updated_node
+
+    def at_least_one_zero_arg(self, original_args: list[cst.Arg]):
+        first_arg = self.resolve_expression(original_args[0].value)
+        second_arg = self.resolve_expression(original_args[1].value)
+        return is_zero(first_arg) or is_zero(second_arg)
+
+
+FixMathIsClose = CoreCodemod(
+    metadata=Metadata(
+        name="fix-math-isclose",
+        summary="Add `abs_tol` to `math.isclose` Call",
+        review_guidance=ReviewGuidance.MERGE_AFTER_REVIEW,
+        references=[
+            Reference(url="https://docs.python.org/3/library/math.html#math.isclose"),
+        ],
+    ),
+    transformer=LibcstTransformerPipeline(FixMathIsCloseTransformer),
+    detector=None,
+)

src/core_codemods/sonar/sonar_fix_math_isclose.py

@@ -0,0 +1,38 @@
+import libcst as cst
+
+from codemodder.codemods.libcst_transformer import LibcstTransformerPipeline
+from codemodder.result import fuzzy_column_match, same_line
+from core_codemods.fix_math_isclose import FixMathIsClose, FixMathIsCloseTransformer
+from core_codemods.sonar.api import SonarCodemod
+
+
+class FixMathIsCloseSonarTransformer(FixMathIsCloseTransformer):
+    def filter_by_result(self, node) -> bool:
+        """
+        Special case result-matching for this rule because the sonar
+        results returned match only the `math.isclose` call without `(...args...)`
+        """
+        match node:
+            case cst.Call():
+                pos_to_match = self.node_position(node)
+                return any(
+                    self.match_location(pos_to_match, result)
+                    for result in self.results or []
+                )
+        return False
+
+    def match_location(self, pos, result):
+        return any(
+            same_line(pos, location) and fuzzy_column_match(pos, location)
+            for location in result.locations
+        )
+
+
+SonarFixMathIsClose = SonarCodemod.from_core_codemod(
+    name="fix-math-isclose-S6727",
+    other=FixMathIsClose,
+    rule_id="python:S6727",
+    rule_name="The abs_tol parameter should be provided when using math.isclose to compare values to 0",
+    rule_url="https://rules.sonarsource.com/python/RSPEC-6727/",
+    transformer=LibcstTransformerPipeline(FixMathIsCloseSonarTransformer),
+)

tests/codemods/sonar/test_sonar_fix_math_isclose.py

@@ -0,0 +1,48 @@
+import json
+
+from codemodder.codemods.test import BaseSASTCodemodTest
+from core_codemods.sonar.sonar_fix_math_isclose import SonarFixMathIsClose
+
+
+class TestSonarFixMathIsClose(BaseSASTCodemodTest):
+    codemod = SonarFixMathIsClose
+    tool = "sonar"
+
+    def test_name(self):
+        assert self.codemod.name == "fix-math-isclose-S6727"
+
+    def test_simple(self, tmpdir):
+        input_code = """
+        import math        
+        
+        def foo(a):
+            return math.isclose(a, 0)
+        """
+        expected_output = """
+        import math        
+        
+        def foo(a):
+            return math.isclose(a, 0, abs_tol=1e-09)
+        """
+        hotspots = {
+            "issues": [
+                {
+                    "rule": "python:S6727",
+                    "status": "OPEN",
+                    "component": "code.py",
+                    "textRange": {
+                        "startLine": 5,
+                        "endLine": 5,
+                        "startOffset": 11,
+                        "endOffset": 23,
+                    },
+                }
+            ]
+        }
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected_output,
+            results=json.dumps(hotspots),
+            num_changes=1,
+        )

tests/codemods/test_fix_math_isclose.py

@@ -0,0 +1,88 @@
+from codemodder.codemods.test import BaseCodemodTest
+from core_codemods.fix_math_isclose import FixMathIsClose
+
+
+class TestFixMathIsClose(BaseCodemodTest):
+    codemod = FixMathIsClose
+
+    def test_name(self):
+        assert self.codemod.name == "fix-math-isclose"
+
+    def test_change(self, tmpdir):
+        input_code = """
+        import math
+        
+        math.isclose(a, 0, abs_tol=0.0)
+        math.isclose(a, 0, abs_tol=0)
+        math.isclose(0, 20)
+
+        def foo(a):
+            return math.isclose(a, 0)
+        
+        """
+        expected = """
+        import math
+        
+        math.isclose(a, 0, abs_tol=1e-09)
+        math.isclose(a, 0, abs_tol=1e-09)
+        math.isclose(0, 20, abs_tol=1e-09)
+
+        def foo(a):
+            return math.isclose(a, 0, abs_tol=1e-09)
+        
+        """
+        self.run_and_assert(tmpdir, input_code, expected, num_changes=4)
+
+    def test_change_resolves_to_zero(self, tmpdir):
+        input_code = """
+        import math
+
+        a = 0
+        math.isclose(a, b)
+
+        c = float(0.0)
+        math.isclose(c, d)
+        math.isclose(1, int(0.0))
+        """
+        expected = """
+        import math
+
+        a = 0
+        math.isclose(a, b, abs_tol=1e-09)
+
+        c = float(0.0)
+        math.isclose(c, d, abs_tol=1e-09)
+        math.isclose(1, int(0.0), abs_tol=1e-09)
+        """
+        self.run_and_assert(tmpdir, input_code, expected, num_changes=3)
+
+    def test_no_change(self, tmpdir):
+        input_code = """
+        import math
+        
+        math.isclose(a, b) 
+        math.isclose(5, 29) 
+        math.isclose(a, 0, abs_tol=0.1)
+        
+        def bar(a, b):
+            return math.isclose(a, b)
+            
+        math.isclose()
+        math.isclose(2)
+        """
+        self.run_and_assert(tmpdir, input_code, input_code)
+
+    def test_exclude_line(self, tmpdir):
+        input_code = (
+            expected
+        ) = """
+        import math
+        math.isclose(20, 0)
+        """
+        lines_to_exclude = [3]
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected,
+            lines_to_exclude=lines_to_exclude,
+        )