Search code examples
pythonmatplotlibpython-class

Select data member in list of custom objects


i have defined the following custom class:

class Point():
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

and I have a list of Point objects called points. I now need to plot this points in a 3D scatter. Is there a quick way to get the x values for all the points that I can implement inside the class definition? I know I can do this with

xs = [p.x for p in points]
ys = ...

but it is a bit tedious. Does anybody know a way to incororate this maybe inside my class? Or maybe I need to define a PointList class?

Thanks


Solution

  • You can create a PointList class.

    class PointList(list):
      def xs(self):
        return [p.x for p in self]
    
      def ys():
        return [p.y for p in self]
    
      def zs():
        return [p.z for p in self]
    

    Then you can use it like this:

    points = PointList([Point(4, 5, 6), Point(2, 6, 4)]) #constructor
    print(points.xs()) # [4, 2]
    

    Since this inherits from list all the standard list methods work.

    Note that these are all functions. You could make them getters that run on each execution but that's added complexity and obscures work that needs to be done. Technically you could store each of these lists and return those lists, but that makes it more complicated to add/remove/update points (and the returned list would be mutable).