How to simplify these two methods?

Question

Well, in my class there are two magical methods, add and sub:

def __add__(self, other):
    if other.__m == self.__m and other.__n == self.__n:
        result = []

        for i in range(other.__m):
            lst = []

            for j in range(other.__n):
                lst.append(self.elems[i][j] + other.elems[i][j])

            result.append(lst)

        return Matrix(result)

def __sub__(self, other):
    if other.__m == self.__m and other.__n == self.__n:
        result = []

        for i in range(other.__m):
            lst = []

            for j in range(other.__n):
                lst.append(self.elems[i][j] - other.elems[i][j])     

            result.append(lst)
        return Matrix(result)

As you can see, the only difference between them is that at a given time one sum and another subtract. Knowing this one difference, I wanted a way to simplify them in order to decrease the code.

Answer 1

You can create a function that accepts the operation (sum, subtraction or other) as argument, and then call this function by modifying only the desired operation.

Since we do not have the direct function, only the operator, and the function varies depending on the type of class of operands, we import operator to provide us with the function related to the correct operation:

import operator

def _fn_base(self, other, op):
    if other.__m == self.__m and other.__n == self.__n:
        result = []

        for i in range(other.__m):
            lst = []

            for j in range(other.__n):
                lst.append(op(self.elems[i][j], other.elems[i][j]))

            result.append(lst)

        return Matrix(result)

def __add__(self, other):
    return self._fn_base(other, operator.add)

def __sub__(self, other):
    return self._fn_base(other, operator.sub)

Answer 2

In addition to creating an auxiliary method that takes the operator as a parameter, as shown in the other answers, you can make use of other magic methods to simplify the logic of your class.

Consider the initial class:

class Matrix:

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

You can make your object eternal by defining the method __iter__:

class Matrix:

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


    def __iter__(self):
        yield from iter(self.matrix)

This way, you can iterate over the matrix lines with the for:

>>> m = Matrix([[1, 2, 3], [4, 5, 6]])

>>> for linha in m:
...     print(linha)

[1, 2, 3]
[4, 5, 6]

You can use the method __len__ to return the dimension of the matrix:

class Matrix:

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


    def __iter__(self):
        yield from iter(self.matrix)


    def __len__(self):
        return len(self.matrix)

And so do:

>>> print(len(m))
2

You can implement the method __getitem__ to facilitate access to matrix lines:

class Matrix:

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


    def __iter__(self):
        yield from iter(self.matrix)


    def __len__(self):
        return len(self.matrix)


    def __getitem__(self, key):
        return self.matrix[key]

Can do:

>>> print(m[0])
[1, 2, 3]

Finally, implementing the addition and subtraction methods, your class could look like:

class Matrix:

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


  @property
  def size(self):
    return (len(self), len(self[0]))


  def __iter__(self):
    yield from iter(self.matrix)


  def __len__(self):
    return len(self.matrix)


  def __getitem__(self, key):
    return self.matrix[key]


  def __add__(self, other):
    if self.size != other.size:
      raise Exception('Matrizes de tamanhos diferentes')

    return Matrix(
      [
        [self[i][j] + other[i][j] for j in range(len(self[i]))] 
          for i in range(len(self))
      ]
    )


  def __sub__(self, other):
    if self.size != other.size:
      raise Exception('Matrizes de tamanhos diferentes')

    return Matrix(
      [
        [self[i][j] - other[i][j] for j in range(len(self[i]))] 
          for i in range(len(self))
      ]
    )

Answer 3

Hey, buddy!
You can include a new parameter in the method and perform the operation according to its value. Example:

def __action__(self, other, mat_operation):
    if other.__m == self.__m and other.__n == self.__n:
        result = []

        for i in range(other.__m):
            lst = []

            for j in range(other.__n):
                if (mat_operation == 1):
                    lst.append(self.elems[i][j] + other.elems[i][j])
                elif(mat_operation == 2):
                    lst.append(self.elems[i][j] - other.elems[i][j])
                else:
                    #another operation


            result.append(lst)

        return Matrix(result)



__action__(other, 1)

I hope I’ve helped.
Hugs,