CODEWITHSUNDEEP
python
user1654183
user1654183

Reputation: 4605

Extending list class to make a list of objects

I have a particle list of objects of type Particle, which takes two parameters, position and energy:

class Particle(object):
    def __init__(self, pos, energy):
        self.x = pos
        self.E = energy

The only way I've managed to do it so far is to create a list of particles using a list comprehension:

number_of_particles = 10
initial_energy = 0
particle_list = [Particle(initial_energy,i) for i in range(number_of_particles)]

which now allows me to do things like:

particle_list[0].x

which is what I want.

However, what I would really, really like is to do something as follows:

particle_list = ParticleList(no_of_particles, initial_energy)

and it create the exact same list.

I assume I have to extend the list class somehow but I'm at a loss as to how to do this.

Upvotes: 2

Views: 86

Answers (3)

DeepSpace
DeepSpace

Reputation: 81594

Create your class with your custom __init__ method.

class ParticleList(list):
    def __init__(self, num, energy):
        self.particle_list = [Particle(energy,i) for i in range(num)]

particles = ParticleList(2, 0).particle_list
for particle in particles:
     print (particle.x, particle.E)

>>(0, 0)
>>(0, 1)

You may create your own method and not use the __init__, this way you will be able to simply return the created list and not assign it to a member (__init__ is not allowed to have a return value).

class ParticleList(list):
    def create_list(self, num, energy):
        return [Particle(energy,i) for i in range(num)]

my_list = ParticleList().create_list(2, 0)

And as others have said, you don't even need the class and can get away with only creating a function:

def create_list(num, energy):
    return [Particle(energy,i) for i in range(num)]

my_list = create_list(2, 0)

Upvotes: 0

kylieCatt
kylieCatt

Reputation: 11039

class Particle(object):
    def __init__(self, pos, energy):
        self.x = pos
        self.E = energy

    @classmethod
    def make_particle_list(self, no_particles, initial_energy=0):
        return [Particle(i, initial_energy) for i in range(no_particles)]

    # this is just for display purposes
    def __repr__(self):
        return 'pos: {p.x} - Energy: {p.E}'.format(p=self)

This offers you a little flexibility. If you only need one particle you can make just one the normal way or:

>>> lst = Particle.make_particle_list(10)
>>> lst
[pos: 0 - Energy: 0, pos: 1 - Energy: 0, pos: 2 - Energy: 0, pos: 3 - Energy: 0, pos: 4 - Energy: 0, pos: 5 - Energy: 0, pos: 6 - Energy: 0, pos: 7 - Energy: 0, pos: 8 - Energy: 0, pos: 9 - Energy: 0]

This also allows you to pass in a different initial_energy if you ever need a different value.

You also had your arguments backwards in your example. You had initial_energy as the first positional argument in your list comprehension but you have it as the second in your __init__() method.

Upvotes: 2

NendoTaka
NendoTaka

Reputation: 1224

Why not just build a function to do this for you. You could do something simple like:

def ParticleList(no_of_particles, initial_energy):
    return [Particle(initial_energy,i) for i in range(number_of_particles)]

This should be a simple way of getting your list.

Upvotes: 2

Related Questions