Cast a Python class to Numpy Array

Tags: , , ,



Can I cast a python class to a numpy array?

from dataclasses import dataclass
import numpy as np
@dataclass
class X:
    x: float = 0
    y: float = 0
x = X()
x_array = np.array(x) # would like to get an numpy array np.array([X.x,X.y])

In the last step, I would like the to obtain an array np.array([X.x, X.y]). Instead, I get array(X(x=0, y=0), dtype=object).

Can I provide method for the dataclass so that the casting works as desired (or overload one of the existing methods of the dataclass)?

Answer

From docstring of numpy.array we can see requirements for the first parameter

object: array_like

An array, any object exposing the array interface, an object whose __array__ method returns an array, or any (nested) sequence.

So we can define a __array__ method like

@dataclass
class X:
    x: float = 0
    y: float = 0

    def __array__(self) -> np.ndarray:
        return np.array([self.x, self.y])

and it will be used by np.array. This should work for any custom Python class I guess.



Source: stackoverflow