Pretty printing of the numpy ndarrays

Question

Following this idea for pretty printing of numpy ndarrays, I have developed a very primitive prototype:

def ndtotext(A, w=None, h=None):
    if A.ndim==1:
        if w == None :
            return str(A)
        else:
            s ='['+' '*(max(w[-1],len(str(A[0])))-len(str(A[0]))) +str(A[0])
            for i,AA in enumerate(A[1:]):
                s += ' '*(max(w[i],len(str(AA)))-len(str(AA))+1)+str(AA)
            s +='] '
    elif A.ndim==2:
        w1 = [max([len(str(s)) for s in A[:,i]])  for i in range(A.shape[1])]
        w0 = sum(w1)+len(w1)+1
        s= u'\u250c'+u'\u2500'*w0+u'\u2510' +'\n'
        for AA in A:
            s += ' ' + ndtotext(AA, w=w1) +'\n'    
        s += u'\u2514'+u'\u2500'*w0+u'\u2518'
    elif A.ndim==3:
        h=A.shape[1]
        s1=u'\u250c' +'\n' + (u'\u2502'+'\n')*h + u'\u2514'+'\n'
        s2=u'\u2510' +'\n' + (u'\u2502'+'\n')*h + u'\u2518'+'\n'
        strings=[ndtotext(a)+'\n' for a in A]
        strings.append(s2)
        strings.insert(0,s1)
        s='\n'.join(''.join(pair) for pair in zip(*map(str.splitlines, strings)))
    return s

for example:

shape = 4, 5, 3
C=np.random.randint(10000, size=np.prod(shape)).reshape(shape)
print(ndtotext(C))

┌┌────────────────┐┌────────────────┐┌────────────────┐┌────────────────┐┐
│ [9298 4404 1759]  [5426 3488 9267]  [8884 7721  579]  [6872 4226 1858] │
│ [6723  271 8466]  [9885 6760 8949]  [ 295 7422 5659]  [5322 4239 7446] │
│ [7156 6077 9390]  [2712 6379 2832]  [6956  626 5534]  [ 142 4090 6390] │
│ [9377 9033 1953]  [8986 3791 4538]  [2466 8572  662]  [1528 8922 9656] │
│ [1449 7319 3939]  [7350 9619  928]  [7542 4704 1477]  [ 980 6037  869] │
└└────────────────┘└────────────────┘└────────────────┘└────────────────┘┘

I would appreciate it if you could review this code and let me know how I can improve it.

I hope to see:

• possible mistakes or cases to break the code
• how to make it faster, more performant, pythonic
• how to extend it to higher dimensions

P.S. For those who follow up this idea I have integrated everythin here in this Jupyter Notebook

Answer 1

If A.ndim is not in 1, 2, 3, your code tries to return a non-existing string s. It would be better to be explicit about what your code supports atm:

def ndtotext(A, w=None, h=None):
    ...
    else:
        raise NotImplementedError("Currently only 1 - 3 dimensions are supported")
    return s

While we are at the point of having your code be clear about what is happening, you should add a docstring explaining what your code does:

def ndtotext(A, w=None, h=None):
    """Returns a string to pretty print the numpy.ndarray `A`.

    Currently supports 1 - 3 dimensions only.
    Raises a NotImplementedError if an array with more dimensions is passed.

    Describe `w` and `h`.
    """
    ...

Next, Python has an official style-guide, PEP8, which programmers are encouraged to follow. One of the things it recommends is surrounding operators with spaces (which I fixed in the rest of the code) and using lower_case for variables and functions (which I left as is for now).

Now, let's come to your actual code:

• You calculate some values multiple times (like str(A[0])), save those to a variable.
• If you want to compare to None, use is (since it is a singleton).
• No else needed after an if...return (this is a matter of personal style, I prefer not having the additional level of indentation).
• Use str.rjust to add enough whitespace in front of your strings. You could also use str.format for this, but it looks less nice.
• Your w has a weird structure, with the width of the first column being the last entry and the rest starting at zero.
• Give the unicode values names. And then add functions to draw a line of specified length.

• String addition is costly and slow. Try to consistently use building a list and str.joining it.

  UPPER_LEFT = u'\u250c'
  UPPER_RIGHT = u'\u2510'
  LOWER_LEFT = u'\u2514'
  LOWER_RIGHT = u'\u2518'
  HORIZONTAL = u'\u2500'
  VERTICAL = u'\u2502'
  
  def upper_line(width):
      return UPPER_LEFT + HORIZONTAL * width + UPPER_RIGHT
  
  def lower_line(width):
      return LOWER_LEFT + HORIZONTAL * width + LOWER_RIGHT
  
  def left_line(height):
      return "\n".join([UPPER_LEFT] + [VERTICAL] * height + [LOWER_LEFT])
  
  def right_line(height):
      return "\n".join([UPPER_RIGHT] + [VERTICAL] * height + [LOWER_RIGHT])
  
  def ndtotext(A, w=None, h=None):
      """Returns a string to pretty print the numpy.ndarray `A`.
  
      Currently supports 1 - 3 dimensions only.
      Raises a NotImplementedError if an array with more dimensions is passed.
  
      Describe `w` and `h`.
      """
      if A.ndim == 1:
          if w is None:
              return str(A)
          s = " ".join([str(value).rjust(width) for value, width in zip(A, w)])
          return '[{}]'.format(s)
      elif A.ndim == 2:
          widths = [max([len(str(s)) for s in A[:, i]]) for i in range(A.shape[1])]
          s = "".join([' ' + ndtotext(AA, w=widths) + ' \n' for AA in A])
          w0 = sum(widths) + len(widths) - 1 + 2 # spaces between elements and corners
          return upper_line(w0) + '\n'  + s + lower_line(w0)
      elif A.ndim == 3:
          h = A.shape[1]
          strings = [left_line(h)]
          strings.extend(ndtotext(a) + '\n' for a in A)
          strings.append(right_line(h))
          return '\n'.join(''.join(pair) for pair in zip(*map(str.splitlines, strings)))
      raise NotImplementedError("Currently only 1 - 3 dimensions are supported")
  

This can probably be even more compactified, but I think it is a good start.

Example usage:

x = np.arange(12)

print(ndtotext(x))
[ 0  1  2  3  4  5  6  7  8  9 10 11]

print(ndtotext(x.reshape(3, 4)))
┌───────────┐
 [0 1  2  3] 
 [4 5  6  7] 
 [8 9 10 11] 
└───────────┘

print(ndtotext(x.reshape(3, 2, 2)))
┌┌─────┐┌─────┐┌───────┐┐
│ [0 1]  [4 5]  [ 8  9] │
│ [2 3]  [6 7]  [10 11] │
└└─────┘└─────┘└───────┘┘