simple magic that enables better use of ipython as a calculator

50-pycalc.py

""" PyCalc

A small "package" that makes it easier for engineers to use IPython as a calculator.

The idea is to extend the python syntax to understand the SI prefixes that are
customary in engineering, like "milli" or "Mega". The code below registers a
TokenInputTransformer with IPython that will rewrite every python line before 
it is given to python itself.

It also changes the IPython output routines so that the same prefixes are also 
used there.

To install just copy the file into your profiles startup folder. Most of the 
time this will be in one of these folders:
    * ~/.ipython/profile_default/startup 
    * ~/.config/ipython/profile_default/startup

The extension is not activated by default, it has to be activated with the
%calc magic.

BEWARE: this will change ipythons output, sometimes e.g. reducing the precision
of a calulated result (because we want to have e.g. 2.34k). Values stored in
variables will have to result with full precision though:

>>> x = 2.333k
>>> x
2.33k
>>> x * 10
23.33k

At the moment numbers are rounded to two decimal digits, not to e.g. a number of
significant digits.
"""

from __future__ import print_function
import math
import sys
from IPython import get_ipython
from IPython.core.inputtransformer import TokenInputTransformer
from IPython.core.magic import register_line_magic
if sys.version_info[0] == 3:
    from IPython.utils._tokenize_py3 import TokenInfo
    conv = lambda t: t
else:
    # ipython running on python2 does not return TokenInfo objects
    # but tuples. To be able to keep the same code, just fake
    # the used functionality of the TokenInfo object
    class TokenInfo(tuple):
        def __new__(cls, *args, **kwargs):
            return super(TokenInfo, cls).__new__(cls, *args, **kwargs)

        def __init__(self, *args, **kwargs):
            super(TokenInfo, self).__init__(*args, **kwargs)
            self.type = self[0]
            self.string = self[1]
            self.start = self[2]
            self.end = self[3]
            self.line = self[4]
    conv = lambda t: TokenInfo(t)
 
@register_line_magic
def calc(line):
    """ activate formatting and input to make use of IPython as calculator easier
    
    Enables SI postfixes:
    >>> 1k/5M
    200m
    
    BEWARE: This will change the output of IPython, you wil e.g. not see floats
    with the complete precision they would have:
    
    >>> 1/3
    0.33
    """
    ip = get_ipython()
    formatter = ip.display_formatter.formatters['text/plain']
    formatter.for_type(int, _intPrettyPrint)
    formatter.for_type(float, _floatPrettyPrint)
 
    for s in (ip.input_splitter, ip.input_transformer_manager):
        s.python_line_transforms.extend([si_postfix_transformer()])
        
def _magnitude(x):
    """ returns the (magnitude of x) / 3 
        only works in the range 1e27 > x > 1e-27, otherwise returns None
    """
    # introduce an offset to that we do not get
    # problems because of the integer division of negative
    # numbers
    n = math.log10(x) + 24
    if math.isinf(n):
        return None
    m = (int(n) // 3) - 8
    # 8 : because 8/-8 is the largest/smallest supported magnitude
    if abs(m) > 8:
        return None
    return m

def _precomma(x):
    """ number of digits before comma when using engineering notation
        only works for numbers x > 1e-101, otherwise returns garbage
    """
    return ((int(math.log10(x) + 100) - 100) % 3) + 1
        
_postfixes = {
    'Y': 1.e24 , 'Z': 1.e21, 'E': 1.e18, 'P': 1.e15, 'T': 1.e12,
    'G': 1.e9, 'M': 1.e6, 'k': 1.e3, '': 1, 'm': 1.e-3, 'u': 1.e-6, 'n': 1.e-9,
    'p': 1.e-12, 'f': 1.e-15, 'a': 1.e-18, 'z': 1.e-21, 'y': 1.e-24,
}
 
_postfixes_rev = { _magnitude(x) : (k,x) for k,x in _postfixes.items() }
 
def _intPrettyPrint(x, p, cycle):
    """ print more information for an integer
    Formatter function that can be registered in IPython to print
    additional information if an integer is being printed to the 
    IPython output.
 
    It will e.g. print
    >>> 5
    5 (0x0005)
    """
 
    if x >= 0:
        p.text('{0} (0x{0:x} uint)'.format(x))
    else:
        xa = abs(x)
        if xa <= 2**15:
            xa = 2**16 - xa
            p.text('{0} (0x{1:F>4x} sint)'.format(x,xa))
        elif xa <= 2**31:
            xa = 2**32 - xa
            p.text('{0} (0x{1:F>8x} sint)'.format(x,xa))
        elif xa <= 2**63:
            xa = 2**64 - xa
            p.text('{0} (0x{1:F>16x} sint)'.format(x,xa))

def _floatPrettyPrint(x, p, cycle):
    """ pretty print floats with SI postfixes
    Formatter function that can be registered in IPython to print
    floats in 'engineering' formatting:
 
    >>> 5e6
    5M
    """
 
    scale = _magnitude(x)
    if scale in _postfixes_rev:
        x /= _postfixes_rev[scale][1]
        p.text('{:.2f}{}'.format(x, _postfixes_rev[scale][0]))
    else:
        p.text('{}'.format(x))
 
@TokenInputTransformer.wrap
def si_postfix_transformer(tokens):
    """ transforms input so that IPython recognizes SI postfixes
 
    >>> 5M
    5e6
    """
 
    result = []
    save = None
    for t in tokens:
        t = conv(t)
        if t.type == 1 and save is not None: # it's a name and we saved a number
            if t.string in _postfixes:
                ns = '{}'.format(float(save.string) * _postfixes[t.string])
                result.append(TokenInfo( 2, ns, save.start, t.end, save.line))
            else:
                result.append(save)
                result.append(t)
            save = None
        elif t.type == 2:
            if save is not None:
                result.append(save)
            save = t
        else:
            if save is not None:
                result.append(save)
                save = None
            result.append(t)
    return result