import numpy

from .stain_color_map import stain_color_map

def get_stain_vector(args, index):
    """Get the stain corresponding to args.stain_$index and
    args.stain_$index_vector.  If the former is not "custom", the
    latter must be None.

    """
    args = vars(args)
    stain = args['stain_' + str(index)]
    stain_vector = args['stain_' + str(index) + '_vector']
    if stain == 'custom':
        if stain_vector is None:
            raise ValueError('If "custom" is chosen for a stain, '
                             'a stain vector must be provided.')
        return stain_vector
    else:
        if stain_vector is None:
            return stain_color_map[stain]
        raise ValueError('Unless "custom" is chosen for a stain, '
                         'no stain vector may be provided.')

def get_stain_matrix(args):
    """Get the stain matrix corresponding to the args.stain_$index and
    args.stain_$index_vector arguments for values of index 1, 2, 3.
    Return a numpy array of column vectors.

    """
    return numpy.array([get_stain_vector(args, i) for i in 1, 2, 3]).T

__all__ = (
    'get_stain_vector',
)

import numpy

from .stain_color_map import stain_color_map

def get_stain_vector(args, index):
    """Get the stain corresponding to args.stain_$index and
    args.stain_$index_vector.  If the former is not "custom", the
    latter must be None.

    """
    args = vars(args)
    stain = args['stain_' + str(index)]
    stain_vector = args['stain_' + str(index) + '_vector']
    if stain == 'custom':
        if stain_vector is None:
            raise ValueError('If "custom" is chosen for a stain, '
                             'a stain vector must be provided.')
        return stain_vector
    else:
        if stain_vector is None:
            return stain_color_map[stain]
        raise ValueError('Unless "custom" is chosen for a stain, '
                         'no stain vector may be provided.')

def get_stain_matrix(args, count=3):
    """Get the stain matrix corresponding to the args.stain_$index and
    args.stain_$index_vector arguments for values of index 1 to count.
    Return a numpy array of column vectors.

    """
    return numpy.array([get_stain_vector(args, i+1) for i in range(count)]).T

__all__ = (
    'get_stain_vector',
)

Enhance get_stain_matrix to take any desired number of vectors
Enhance get_stain_matrix to take any desired number of vectors

  import numpy
  
  from .stain_color_map import stain_color_map
  
  def get_stain_vector(args, index):
      """Get the stain corresponding to args.stain_$index and
      args.stain_$index_vector.  If the former is not "custom", the
      latter must be None.
  
      """
      args = vars(args)
      stain = args['stain_' + str(index)]
      stain_vector = args['stain_' + str(index) + '_vector']
      if stain == 'custom':
          if stain_vector is None:
              raise ValueError('If "custom" is chosen for a stain, '
                               'a stain vector must be provided.')
          return stain_vector
      else:
          if stain_vector is None:
              return stain_color_map[stain]
          raise ValueError('Unless "custom" is chosen for a stain, '
                           'no stain vector may be provided.')
  
- def get_stain_matrix(args):
+ def get_stain_matrix(args, count=3):
      """Get the stain matrix corresponding to the args.stain_$index and
-     args.stain_$index_vector arguments for values of index 1, 2, 3.
+     args.stain_$index_vector arguments for values of index 1 to count.
      Return a numpy array of column vectors.
  
      """
-     return numpy.array([get_stain_vector(args, i) for i in 1, 2, 3]).T
+     return numpy.array([get_stain_vector(args, i+1) for i in range(count)]).T
  
  __all__ = (
      'get_stain_vector',
  )
  
Enhance get_stain_matrix to take any desired number of vectors
## Code Before:
import numpy

from .stain_color_map import stain_color_map

def get_stain_vector(args, index):
    """Get the stain corresponding to args.stain_$index and
    args.stain_$index_vector.  If the former is not "custom", the
    latter must be None.

    """
    args = vars(args)
    stain = args['stain_' + str(index)]
    stain_vector = args['stain_' + str(index) + '_vector']
    if stain == 'custom':
        if stain_vector is None:
            raise ValueError('If "custom" is chosen for a stain, '
                             'a stain vector must be provided.')
        return stain_vector
    else:
        if stain_vector is None:
            return stain_color_map[stain]
        raise ValueError('Unless "custom" is chosen for a stain, '
                         'no stain vector may be provided.')

def get_stain_matrix(args):
    """Get the stain matrix corresponding to the args.stain_$index and
    args.stain_$index_vector arguments for values of index 1, 2, 3.
    Return a numpy array of column vectors.

    """
    return numpy.array([get_stain_vector(args, i) for i in 1, 2, 3]).T

__all__ = (
    'get_stain_vector',
)

## Instruction:
Enhance get_stain_matrix to take any desired number of vectors
## Code After:
import numpy

from .stain_color_map import stain_color_map

def get_stain_vector(args, index):
    """Get the stain corresponding to args.stain_$index and
    args.stain_$index_vector.  If the former is not "custom", the
    latter must be None.

    """
    args = vars(args)
    stain = args['stain_' + str(index)]
    stain_vector = args['stain_' + str(index) + '_vector']
    if stain == 'custom':
        if stain_vector is None:
            raise ValueError('If "custom" is chosen for a stain, '
                             'a stain vector must be provided.')
        return stain_vector
    else:
        if stain_vector is None:
            return stain_color_map[stain]
        raise ValueError('Unless "custom" is chosen for a stain, '
                         'no stain vector may be provided.')

def get_stain_matrix(args, count=3):
    """Get the stain matrix corresponding to the args.stain_$index and
    args.stain_$index_vector arguments for values of index 1 to count.
    Return a numpy array of column vectors.

    """
    return numpy.array([get_stain_vector(args, i+1) for i in range(count)]).T

__all__ = (
    'get_stain_vector',
)
