Cast a Python class to Numpy Array

Viewed 1154

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)?

1 Answers

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 an __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.

Related