Source code for irfpy.util.with_context

''' With statement extension.

With `with statement <https://docs.python.org/3/reference/compound_stmts.html#with>`_, a code block is wrapperd
by a `context manager <https://docs.python.org/3/reference/datamodel.html#context-managers>`_.
See also https://docs.python.org/3/library/contextlib.html

.. codeauthor:: Yoshifumi Futaana

.. autosummary::
	
    printoptions
'''

import contextlib

import numpy as np

@contextlib.contextmanager
def printoptions(*args, **kwds):
    ''' This with statement controls the output of numpy array.
    
    Use numpy.set_printoptions locally in with-block.
    
    :params args: Identical to ``numpy.set_printoptions``.
    :keywords kwds: Identical to ``numpy.set_printoptions``.
     
    Simple usage is
    
    >>> with printoptions(precision=3, suppress=True):
    ...     print(np.logspace(-3, 3, 5))    # doctest: +NORMALIZE_WHITESPACE
    [    0.001     0.032     1.       31.623  1000.   ]

    Step-by-step explanation follows.
    
    >>> x = np.logspace(-3, 3, 5)

    A default output will be

    >>> print(x)   # doctest: +NORMALIZE_WHITESPACE
    [1.00000000e-03 3.16227766e-02 1.00000000e+00 3.16227766e+01
     1.00000000e+03]

    Such output has too much digits, and thus redundunt.
    Also, sometimes you want to suppress insignificant value
    like 3.8132234e-17 as 0.
    
    Numpy prepares the functionality, ``set_printoptions``,
    but it will change the numpy feature globally.

    Therefore, use :func:`printoptions`.
    Just do as follows.

    >>> with printoptions(precision=3, suppress=True):
    ...     print(x)    # doctest: +NORMALIZE_WHITESPACE
    [    0.001     0.032     1.       31.623  1000.   ]

    Outside with statement, the print option becomes back.

    >>> print(x)  # doctest: +NORMALIZE_WHITESPACE
    [1.00000000e-03   3.16227766e-02   1.00000000e+00   3.16227766e+01
     1.00000000e+03]

    The code is taken from
    http://stackoverflow.com/questions/2891790/pretty-printing-of-numpy-array
    '''
    original = np.get_printoptions()
    np.set_printoptions(*args, **kwds)
    yield
    np.set_printoptions(**original)


import unittest
import doctest


def doctests():
    return unittest.TestSuite((
        doctest.DocTestSuite(),
        ))
if __name__ == '__main__':
    unittest.main(defaultTest='doctests')