<s> import numpy as np <EOL> from chaco . api import AbstractPlotData , ArrayPlotData , Plot , ArrayDataSource <EOL> from traits . api import Dict , Instance , Str <EOL> from pandas import DataFrame <EOL> class PandasPlotData ( AbstractPlotData ) : <EOL> '''<STR_LIT>''' <EOL> df = Instance ( DataFrame ) <EOL> extras = Dict ( ) <EOL> selections = Dict ( ) <EOL> def __init__ ( self , dataframe , columnlabel = '<STR_LIT>' , indexlabel = '<STR_LIT:index>' ) : <EOL> """<STR_LIT>""" <EOL> super ( AbstractPlotData , self ) . __init__ ( ) <EOL> if len ( dataframe . shape ) > <NUM_LIT:2> : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> self . columnlabel = columnlabel <EOL> self . indexlabel = indexlabel <EOL> event = self . _add_dataframe ( dataframe ) <EOL> self . data_changed = event <EOL> def list_data ( self , axis = <NUM_LIT:0> , as_strings = False ) : <EOL> """<STR_LIT>""" <EOL> if axis == <NUM_LIT:0> or axis == self . columnlabel : <EOL> if as_strings : <EOL> datalist = self . _colmasks . keys ( ) <EOL> else : <EOL> datalist = self . _colmasks . values ( ) <EOL> elif axis == <NUM_LIT:1> or axis == self . indexlabel : <EOL> if as_strings : <EOL> datalist = self . _rowmasks . keys ( ) <EOL> else : <EOL> datalist = self . _rowmasks . values ( ) <EOL> return datalist <EOL> def get_data ( self , name ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> return self . df [ name ] . values <EOL> except KeyError : <EOL> try : <EOL> return self . extras [ name ] <EOL> except KeyError : <EOL> try : <EOL> return self . df [ self . _colmasks [ name ] ] . values <EOL> except KeyError : <EOL> try : <EOL> self . extras [ self . _colmasks [ name ] ] . values <EOL> except KeyError : <EOL> if name == self . columnlabel : <EOL> return self . _colmasks . values ( ) <EOL> elif name == self . indexlabel : <EOL> return self . _rowmasks . values ( ) <EOL> def get_row_data ( self , name ) : <EOL> '''<STR_LIT>''' <EOL> return self . df . xs ( name ) <EOL> def set_data ( self , name , new_data ) : <EOL> """<STR_LIT>""" <EOL> if not self . writable : <EOL> return None <EOL> if isinstance ( new_data , list ) or isinstance ( new_data , tuple ) : <EOL> new_data = np . array ( new_data ) <EOL> event = { } <EOL> try : <EOL> self . df [ name ] = new_data <EOL> event [ '<STR_LIT>' ] = [ str ( name ) ] <EOL> except KeyError : <EOL> self . df [ self . _colmasks [ name ] ] = new_data <EOL> event [ '<STR_LIT>' ] = [ str ( name ) ] <EOL> self . data_changed = event <EOL> return name <EOL> def update_dataframe ( self , dataframe ) : <EOL> '''<STR_LIT>''' <EOL> for column in dataframe . columns . values : <EOL> self . set_data ( column , dataframe [ column ] ) <EOL> def set_dataframe ( self , dataframe ) : <EOL> '''<STR_LIT>''' <EOL> removed = { '<STR_LIT>' : self . _colmasks . keys ( ) } <EOL> added = self . _add_dataframe ( dataframe ) <EOL> event = dict ( added . items ( ) + removed . items ( ) ) <EOL> self . data_changed = event <EOL> def _add_dataframe ( self , dataframe ) : <EOL> '''<STR_LIT>''' <EOL> self . df = dataframe <EOL> self . _colmasks = dict ( ( ( str ( val ) , val ) for val in self . df . columns . values ) ) <EOL> self . _rowmasks = dict ( ( ( str ( val ) , val ) for val in self . df . index . values ) ) <EOL> return { '<STR_LIT>' : self . _colmasks . keys ( ) } <EOL> def get_selection ( self , name ) : <EOL> """<STR_LIT>""" <EOL> return self . selections . get ( name , None ) <EOL> def set_selection ( self , name , selection ) : <EOL> self . selections [ name ] = selection <EOL> self . data_changed = True <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> df = DataFrame ( ( np . random . randn ( <NUM_LIT:10> , <NUM_LIT:10> ) ) ) <EOL> data = PandasPlotData ( df ) <EOL> print data . df [ <NUM_LIT:0> ] <EOL> df2 = DataFrame ( ( np . random . randn ( <NUM_LIT:10> , <NUM_LIT:10> ) ) ) <EOL> data . set_dataframe ( df2 ) <EOL> print data . df [ <NUM_LIT:0> ] </s>
