Python Square Root for Class Instances

Question

I am currently implementing a class that can handle numeric data associated with physical units.

I would like to implement a way of calculating the square root of an instance. Assume that you have an instance of a class that has attributes value and name:

from math import sqrt

class Foo:
   def __init__(self, value, name)
      self.value = value
      self.name = name

   def __sqrt__(self):
      return sqrt(self.value)

I would like to implement a function similar to the magic methods like add(self, other) that would calculate the squareroot when I call the math.sqrt() function:

A = Foo(4, "meter")
root = math.sqrt(A)

should return call the A.sqrt() function.

Answer 1

You can't without reassigning math.sqrt to a custom function. If you want to allow Foo to be cast to int or float you can implement __int__ and __float__ and cast before calling math.sqrt:

class Foo:
    def __init__(self, value, name)
        self.value = value
        self.name = name

    def __float__(self):
        return float(self.value)

    def __int__(self):
        return int(self.value)

A = Foo(4, "meter")
root = math.sqrt(float(A))

EDIT: According to the comments below, it seems that you can invoke math.sqrt(A) directly if Foo implements __float__ due to how the math module is implemented. I would still rather be explicit than implicit.

Answer 2

There's no trivial/out-of-the-box way to wire math.sqrt to call Foo.__sqrt__.

Just implement sqrt in Foo:

class Foo:
     ...
     def sqrt(self):
          return sqrt(self.value)

A = Foo(4, "meter")
root = A.sqrt()

If for some reason you insist, it can be hacked, but I don't see any reason why you would want to do that:

from math import sqrt

class Foo:
    def __init__(self, value, name):
        self.value = value
        self.name = name

    def __sqrt__(self):
        return sqrt(self.value)

orig_sqrt = sqrt

def my_sqrt(value):
    if isinstance(value, Foo):
        return orig_sqrt(value.value)
        # or return value.__sqrt__()
    else:
        return orig_sqrt(value)

sqrt = my_sqrt

A = Foo(4, "meter")
print(sqrt(A))
# 2.0
print(sqrt(4))
# 2.0