att/data.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
Data module for a Kai's Text Tools.

(c) 2013 Ivan "Kai SD" Korystin 

License: GPLv3
'''

import csv, codecs

class Data(object):
    '''
    Empty data class. Can be used for a subclassing or procedural data creation.
    '''
    def __init__(self, *args, **kwargs):
        '''
        Constructor
        '''
        self.keys = []
        self.rows = []
    
    def __getitem__(self, pair):
        '''
        Returns a value for given key and row. 
        '''
        key = pair[0]
        row = pair[1]
        
        keys = self.keys
        rows = self.rows
        if key in keys:
            if len(rows) > row:
                return rows[row][keys.index(key)]
            else:
                raise BaseException('Row %i not found in data' % (row))
        else:
            raise BaseException('Named value %s not found in data' % (key))
    
    def __setitem__(self, pair, value):
        '''
        Sets a value for given key and row.  
        '''
        key = pair[0]
        row = pair[1]
        
        keys = self.keys
        rows = self.rows
        if key in keys:
            if len(rows) > row:
                rows[row][keys.index(key)] = value
            else:
                raise BaseException('Row %i not found in data' % (row))
        else:
            raise BaseException('Named value %s not found in data' % (key))
    
    def __str__(self):
        '''
        Returns data as string.
        '''
        return str((self.keys, self.rows))
    
    def __repr__(self):
        return self.__str__()
    
    def has_key(self, key):
        '''
        Returns True if given key exists in data
        '''
        return key in self.keys
    
    def add_rows(self, n=1):
        '''
        Adds some empty rows to the data.
        '''
        keys = self.keys
        rows = self.rows
        
        for n in xrange(0, n):
            row = []
            for k in keys:
                row.append('')
            rows.append(row)
    
    def add_keys(self, *h):
        '''
        Adds new keys to the data.
        '''
        keys = self.keys
        rows = self.rows
        
        for i in h:
            keys.append(i)
        for r in rows:
            for i in h:
                r.append('')
    
    def col_by_key(self, key):
        cols = []
        keys = self.keys
        rows = self.rows
        if key in keys:
            idx = keys.index(key)
            for r in rows:
                cols.append(r[idx])
        else:
            raise BaseException('Named value %s not found in data' % (key))
        return tuple(cols)
    
    def row_by_idx(self, idx):
        return tuple(self.rows[idx])

class CSVData(Data):
    '''
    Class for reading CSV files.
    '''
    class Reader:
        class Recoder:
            def __init__(self, f, encoding):
                self.reader = codecs.getreader(encoding)(f)
        
            def __iter__(self):
                return self
        
            def next(self):
                return self.reader.next().encode("utf-8")
        
        def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
            f = self.Recoder(f, encoding)
            self.reader = csv.reader(f, dialect=dialect, **kwds)
    
        def next(self):
            row = self.reader.next()
            return [unicode(s, "utf-8") for s in row]
    
        def __iter__(self):
            return self
    
    def __init__(self, filename, encoding='utf-8', delimiter=';', quotechar='"', **kwargs):
        csvfile = self.Reader(open(filename), encoding=encoding, delimiter=delimiter, quotechar=quotechar)
        sourceData = []
        sourcekeys = None
        
        if kwargs.get('transpose', False):
            sourcekeys = []
            rowData = []
            for i in csvfile:
                sourcekeys.append(i[0])
                for k in xrange(1, len(i)):
                        sourceData.append([])
                        try:
                            i[k] = int(i[k])
                        except:
                            try:
                                i[k] = float(i[k])
                            except:
                                i[k] = i[k]
                rowData.append(i[1:])
            sourceData = list(map(lambda *x:x, *rowData))
        else:
            for i in csvfile:
                if sourcekeys is None:
                    sourcekeys = i
                else:
                    for k in xrange(0, len(i)):
                            try:
                                i[k] = int(i[k])
                            except:
                                try:
                                    i[k] = float(i[k])
                                except:
                                    i[k] = i[k]
                    sourceData.append(i)
        self.keys = sourcekeys
        self.rows = sourceData