How to Implement pytest.approx() for Data Classes

Viewed 648

Suppose that I have a Python Data Class that I'd like to test with pytest:

@dataclass
class ExamplePoint:
    x: float
    y: float

With simple mathematical operations all is well:

p1 = ExamplePoint(1,2)
p2 = ExamplePoint(0.5 + 0.5, 2)
p1 == p2   # True

But floating point arithmetic quickly causes problems:

p3 = ExamplePoint(1, math.sqrt(2) * math.sqrt(2))
p1 == p3  #  False

Using pytest you can get round this with the approx() function:

2.0 == approx(math.sqrt(2) * math.sqrt(2))  # True

You can't simplemindedly extend this to the Data Class:

p1 = approx(p3) # Results error: TypeError: cannot make approximate comparisons to non-numeric values: ExamplePoint(x=1, y=2.0000000000000004) 

My current solution to fix this is to write an approx() function on the Data Class as:

from dataclasses import astuple, dataclass
import pytest

@dataclass
class ExamplePoint:
    x: float
    y: float
    
    def approx(self, other):
        return astuple(self) == pytest.approx(astuple(other))
    
p1 = ExamplePoint(1,2)
p3 = ExamplePoint(1, math.sqrt(2) * math.sqrt(2))
p1.approx(p3)  # True

I don't like this solution as ExamplePoint now depends on pytest. That seems wrong.

How can I extend pytest so that approx() works with my Data Class without having the Data Class know about pytest?

3 Answers

You can use math.isclose(), which also allows you to set a tolerance for what you consider close. In the following code, I have applied it separately to x and y coordinates of your ExamplePoint outside your dataclass, but you can implement it differently:

@dataclass
class ExamplePoint:
    x: float
    y: float

p1 = ExamplePoint(1, 2)
p3 = ExamplePoint(1, math.sqrt(2) * math.sqrt(2))

print(math.isclose(p1.x, p3.x, rel_tol=0.01)) #True
print(math.isclose(p1.y, p3.y, rel_tol=0.01)) #True

UPDATE: Here is how you can incorporate it into your approx function:

from dataclasses import astuple, dataclass
import math
@dataclass
class ExamplePoint:
    x: float
    y: float

    def approx(self, other):
        return math.isclose(self.x,other.x, rel_tol=0.001) \
               and math.isclose(self.y,other.y, rel_tol=0.001)


p1 = ExamplePoint(1, 2)
p2 = ExamplePoint(1,1.99)
p3 = ExamplePoint(1, math.sqrt(2) * math.sqrt(2))

print(p1.approx(p2)) #False
print(p1.approx(p3)) #True

Digging into the pytest code (see https://github.com/pytest-dev/pytest/blob/main/src/_pytest/python_api.py ), it appears that pytest checks whether the expected value is Iterable and Sizeable. I can make my Data Class Iterable and Sizeable as follows.

from dataclasses import astuple, dataclass
import pytest
import math

@dataclass
class ExamplePoint:
    x: float
    y: float

    def approx(self, other):
        return astuple(self) == pytest.approx(astuple(other))

    def __iter__(self):
        return iter(astuple(self))

    def __len__(self):
        return len(astuple(self))

p1 = ExamplePoint(1,2)
p3 = ExamplePoint(1, math.sqrt(2) * math.sqrt(2))

assert p1 == pytest.approx(p3)  # True
Related