Beginnings of a neuron simulation

Question

I have 2 classes that do extremely similar things. The only differences are that they draw using different colors, and that one draws in the back.

The classes NeuronOutput and NeuronReciever are very similar.

class NeuronOutput():
    """The place where a Neuron outputs its signal"""
    def __init__(self,neuron):
        self.neuron=neuron

        width=self.neuron.width*0.5
        offset=width/2
        distance=self.neuron.width*0.4
        self.pos=[self.neuron.x+math.sin(self.neuron.angle)*distance,self.neuron.y-math.cos(self.neuron.angle)*distance]
        self.drawRect=[self.pos[0]-offset,self.pos[1]-offset,width,width]        
    def draw(self,surf):
        pygame.draw.ellipse(surf,[200,150,200],self.drawRect)
class NeuronReciever():
    """Where a Neuron recieves input"""
    def __init__(self,neuron):
        self.neuron=neuron

        width=self.neuron.width*0.5
        offset=width/2
        distance=self.neuron.width*0.4
        self.pos=[self.neuron.x-math.sin(self.neuron.angle)*distance,self.neuron.y+math.cos(self.neuron.angle)*distance]
        self.drawRect=[self.pos[0]-offset,self.pos[1]-offset,width,width]        
    def draw(self,surf):
        pygame.draw.ellipse(surf,[200,255,200],self.drawRect)

There is also a class Neuron that uses instances of the two above classes as properties. It is also fairly similar to them, however it is a bit different in the way it initializes and draws.

class Neuron():
    """A single neuron with Synapses that connect to others"""
    def __init__(self,x,y):
        self.x=x
        self.y=y
        self.pos=[x,y]
        self.width=30
        self.drawOffset=self.width/2;
        self.connections=[]
        self.angle=math.radians(random.randint(0,360))
        self.drawRect=[self.x-self.drawOffset,self.y-self.drawOffset,self.width,self.width]
        self.output=NeuronOutput(self)
        self.input=NeuronReciever(self)

    def draw(self,surf):

        pygame.draw.ellipse(surf,[200,150,50],self.drawRect)
        self.output.draw(surf)
        self.input.draw(surf)
        self.syn.draw(surf)

(The code is the beginning of a neuron simulation.)

I want to keep all these as separate classes for the sake of readability and ease of use. How can I reduce the repetitiveness of this code, while still keeping it readable?

Answer 1

So you have two classes which are very similar and differ only in one thing? Then make a parent class that has all the behavior and let the others inherit from it or add a switch to flip between the two behaviors.

I would probably go with the first. To accomplish this, make self.drawRect a property (and rename it to self.rect, it does not actually draw anything, it just says where to draw):

class NeuronIO:
    """Draw input or output of a neuron.

    Abstract base class.
    """
    color = 0, 0, 0  # default color is black

    def __init__(self, neuron):
        self.neuron = neuron

        width = 0.5 * self.neuron.width
        offset = width / 2
        distance = 0.4 * self.neuron.width
        self._pos = None
        self._rect = None

    def draw(self, surf):
        pygame.draw.ellipse(surf, self.color, self.rect)

    @property
    def pos(self):
        """Needs to be implemented by the child classes"""
        raise NotImplementedError

    @property
    def rect(self):
        if self._rect is None:
            pos = self.pos
            self._rect = pos[0] - offset, pos[1] - offset, width, width 
        return self._rect


class NeuronOutput(NeuronIO):
    """The place where a Neuron outputs its signal"""
    color = 200, 150, 200

    @property
    def pos(self):
        if self._pos is None:
            self._pos = (self.neuron.x + math.sin(self.neuron.angle) * distance,
                         self.neuron.y - math.cos(self.neuron.angle) * distance)
        return self._pos


class NeuronReciever(NeuronIO):
    """Where a Neuron receives input"""
    color = 200, 255, 200
    @property
    def pos(self):
        if self._pos is None:
            self._pos = (self.neuron.x - math.sin(self.neuron.angle) * distance,
                         self.neuron.y + math.cos(self.neuron.angle) * distance)
        return self._pos

Here both rect and pos are cached so that they are only calculated the first time they are needed.

In addition I added some whitespace around operators and after commas, as recommended by Python's official style-guide, PEP8.

---

At this point you should ask yourself: "But how often do I need to draw the receiver or output without also drawing a neuron?". If the answer to that question is "probably never!", then you should just fold all of that functionality into the Neuron class:

import math
import random
import pygame


class Neuron:
    """A single neuron with Synapses that connect to others"""

    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.pos = x, y
        self.width = 30
        self.offset = self.width / 2
        self.connections = []
        self.angle = math.radians(random.randint(0, 360))
        self.rect = (self.x - self.offset, self.y - self.offset,
                     self.width, self.width)
        self.color = 200, 150, 50

    def draw(self, surf):
        pygame.draw.ellipse(surf, self.color, self.rect)
        pygame.draw.ellipse(surf, *self.input)
        pygame.draw.ellipse(surf, *self.output)

    @property
    def input(self):
        color = 200, 255, 200
        rect = self.calc_rect()
        return color, rect

    @property
    def output(self):
        color = 200, 255, 200
        rect = self.calc_rect(output=True)
        return color, rect

    def calc_rect(self, output=False):
        width = 0.5 * self.width
        distance = 0.4 * self.width
        offset = 0.25 * self.width
        if output:
            pos = (self.x + math.sin(self.angle) * distance,
                   self.y - math.cos(self.angle) * distance)
        else:
            pos = (self.x - math.sin(self.angle) * distance,
                   self.y + math.cos(self.angle) * distance)
        return pos[0] - offset, pos[1] - offset, width, width

Here I did not go through the effort of making sure that Neuron.input and Neuron.output are only calculated once, but you could do it in the same way as I did above.

Note that I made all of your lists into tuples (like Neuron.pos). This is just aesthetics atm, but if you want to be able to do e.g. neuron.pos[0] += 1, then you need to make them lists again.