whos.xsh

"""Namespace related functionality for the xonsh shell.

The following whos routine was forked from the IPython project:
* Copyright (c) 2008-2014, IPython Development Team
* Copyright (C) 2001-2007 Fernando Perez <[email protected]>
* Copyright (c) 2001, Janko Hauser <[email protected]>
* Copyright (c) 2001, Nathaniel Gray <[email protected]>
"""
from xonsh.built_ins import XSH
from xonsh.tools import unthreadable
import sys

START_NAMES = frozenset()


@XSH.builtins.events.on_pre_cmdloop
def _store_start_names():
    global START_NAMES
    START_NAMES = frozenset(XSH.ctx)


def _get_user_keys(ctx):
    def is_valid(name):
        return name not in START_NAMES and not name.startswith("_")
    return [n for n in ctx.keys() if is_valid(n)]


def _filter_by_type(ctx, keys, type_names=()):
    if not type_names:
        return keys
    type_names = set(type_names)
    return [n for n in keys if type(ctx[n]).__name__ in type_names]


def _get_namespace(ctx, type_names=()):
    names = _get_user_keys(ctx)
    filtered = _filter_by_type(ctx, names, type_names)
    return names, filtered
    

@unthreadable
def _whox(args):
    names, filtered = _get_namespace(XSH.ctx, args)
    
    if not filtered:
        if not names:
            return "Interactive namespace is empty."
        return "No variables match your requested type."
        
    return "\t".join(filtered)


@unthreadable
def _whos(args, stdin=None, stdout=None):
    nanes, varnames = _get_namespace(XSH.ctx, args)
    
    if not varnames:
        if not names:
            return "Interactive namespace is empty."
        return "No variables match your requested type."

    seq_types = ["dict", "tuple", "list"]


    ndarray_type = None
    if 'numpy' in sys.modules:
        try:
            from numpy import ndarray
        except ImportError:
            pass
        else:
            ndarray_type = ndarray.__name__
    
    abbrevs = {'IPython.core.macro.Macro' : 'Macro'}
    def type_name(v):
        tn = type(v).__name__
        return abbrevs.get(tn,tn)

    varlist = [XSH.ctx[n] for n in varnames]

    typelist = []
    for vv in varlist:
        tt = type_name(vv)

        if tt=='instance':
            typelist.append( abbrevs.get(str(vv.__class__),
                                         str(vv.__class__)))
        else:
            typelist.append(tt)

    # column labels and # of spaces as separator
    varlabel = 'Variable'
    typelabel = 'Type'
    datalabel = 'Data/Info'
    colsep = 3
    # variable format strings
    vformat    = "{0:<{varwidth}}{1:<{typewidth}}"
    aformat    = "%s: %s elems, type `%s`, %s bytes"
    # find the size of the columns to format the output nicely
    varwidth = max(max(map(len,varnames)), len(varlabel)) + colsep
    typewidth = max(max(map(len,typelist)), len(typelabel)) + colsep
    # table header
    print(varlabel.ljust(varwidth) + typelabel.ljust(typewidth) + \
          ' '+datalabel+'\n' + '-'*(varwidth+typewidth+len(datalabel)+1), file=stdout)
    # and the table itself
    kb = 1024
    Mb = 1048576  # kb**2
    for vname,var,vtype in zip(varnames,varlist,typelist):
        print(vformat.format(vname, vtype, varwidth=varwidth, typewidth=typewidth), end=' ', file=stdout)
        if vtype in seq_types:
            print("n="+str(len(var)), file=stdout)
        elif vtype == ndarray_type:
            vshape = str(var.shape).replace(',','').replace(' ','x')[1:-1]
            if vtype==ndarray_type:
                # numpy
                vsize  = var.size
                vbytes = vsize*var.itemsize
                vdtype = var.dtype

            if vbytes < 100000:
                print(aformat % (vshape, vsize, vdtype, vbytes), file=stdout)
            else:
                print(aformat % (vshape, vsize, vdtype, vbytes), end=' ', file=stdout)
                if vbytes < Mb:
                    print('(%s kb)' % (vbytes/kb,), file=stdout)
                else:
                    print('(%s Mb)' % (vbytes/Mb,), file=stdout)
        else:
            try:
                vstr = str(var)
            except UnicodeEncodeError:
                vstr = var.encode(DEFAULT_ENCODING,
                                  'backslashreplace')
            except:
                vstr = "<object with id %d (str() failed)>" % id(var)
            vstr = vstr.replace('\n', '\\n')
            if len(vstr) < 50:
                print(vstr, file=stdout)
            else:
                print(vstr[:25] + "<...>" + vstr[-25:], file=stdout)


XSH.aliases['whox'] = _whox
XSH.aliases['whos'] = _whos